.NET: Add WithAITool extensions for Hosting AIAgents (#1990)

* add extensions to register tools via fluentAPI

* fix

dotnet/samples/AgentWebChat/AgentWebChat.AgentHost/AgentWebChat.AgentHost.csproj

@@ -1,4 +1,4 @@
-<Project Sdk="Microsoft.NET.Sdk.Web">
+<Project Sdk="Microsoft.NET.Sdk.Web">
 
   <PropertyGroup>
     <TargetFramework>net9.0</TargetFramework>
@@ -13,7 +13,7 @@
     <ProjectReference Include="..\..\..\src\Microsoft.Agents.AI.Hosting\Microsoft.Agents.AI.Hosting.csproj" />
     <ProjectReference Include="..\..\..\src\Microsoft.Agents.AI\Microsoft.Agents.AI.csproj" />
     <ProjectReference Include="..\..\..\src\Microsoft.Agents.AI.Hosting.A2A.AspNetCore\Microsoft.Agents.AI.Hosting.A2A.AspNetCore.csproj" />
-	  <ProjectReference Include="..\..\..\src\Microsoft.Agents.AI.Hosting.OpenAI\Microsoft.Agents.AI.Hosting.OpenAI.csproj" />
+    <ProjectReference Include="..\..\..\src\Microsoft.Agents.AI.Hosting.OpenAI\Microsoft.Agents.AI.Hosting.OpenAI.csproj" />
     <ProjectReference Include="..\AgentWebChat.ServiceDefaults\AgentWebChat.ServiceDefaults.csproj" />
   </ItemGroup>
 
@@ -37,4 +37,4 @@
   </ItemGroup>
   <!-- A2A dependency -->
 
-</Project>
+</Project>

dotnet/samples/AgentWebChat/AgentWebChat.AgentHost/Custom/CustomAITools.cs

@@ -0,0 +1,17 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using Microsoft.Extensions.AI;
+
+namespace AgentWebChat.AgentHost.Custom;
+
+public class CustomAITool : AITool
+{
+}
+
+public class CustomFunctionTool : AIFunction
+{
+    protected override ValueTask<object?> InvokeCoreAsync(AIFunctionArguments arguments, CancellationToken cancellationToken)
+    {
+        return new ValueTask<object?>(arguments.Context?.Count ?? 0);
+    }
+}

dotnet/samples/AgentWebChat/AgentWebChat.AgentHost/Program.cs

@@ -2,6 +2,7 @@
 
 using A2A.AspNetCore;
 using AgentWebChat.AgentHost;
+using AgentWebChat.AgentHost.Custom;
 using AgentWebChat.AgentHost.Utilities;
 using Microsoft.Agents.AI;
 using Microsoft.Agents.AI.Hosting;
@@ -25,6 +26,8 @@ var pirateAgentBuilder = builder.AddAIAgent(
     instructions: "You are a pirate. Speak like a pirate",
     description: "An agent that speaks like a pirate.",
     chatClientServiceKey: "chat-model")
+    .WithAITool(new CustomAITool())
+    .WithAITool(new CustomFunctionTool())
     .WithInMemoryThreadStore();
 
 var knightsKnavesAgentBuilder = builder.AddAIAgent("knights-and-knaves", (sp, key) =>

dotnet/src/Microsoft.Agents.AI.Hosting/AgentHostingServiceCollectionExtensions.cs

@@ -1,6 +1,7 @@
 // Copyright (c) Microsoft. All rights reserved.
 
 using System;
+using System.Collections.Generic;
 using System.Linq;
 using Microsoft.Agents.AI.Hosting.Local;
 using Microsoft.Extensions.AI;
@@ -29,7 +30,8 @@ public static class AgentHostingServiceCollectionExtensions
         return services.AddAIAgent(name, (sp, key) =>
         {
             var chatClient = sp.GetRequiredService<IChatClient>();
-            return new ChatClientAgent(chatClient, instructions, key);
+            var tools = GetRegisteredToolsForAgent(sp, name);
+            return new ChatClientAgent(chatClient, instructions, key, tools: tools);
         });
     }
 
@@ -46,7 +48,11 @@ public static class AgentHostingServiceCollectionExtensions
     {
         Throw.IfNull(services);
         Throw.IfNullOrEmpty(name);
-        return services.AddAIAgent(name, (sp, key) => new ChatClientAgent(chatClient, instructions, key));
+        return services.AddAIAgent(name, (sp, key) =>
+        {
+            var tools = GetRegisteredToolsForAgent(sp, name);
+            return new ChatClientAgent(chatClient, instructions, key, tools: tools);
+        });
     }
 
     /// <summary>
@@ -65,7 +71,8 @@ public static class AgentHostingServiceCollectionExtensions
         return services.AddAIAgent(name, (sp, key) =>
         {
             var chatClient = chatClientServiceKey is null ? sp.GetRequiredService<IChatClient>() : sp.GetRequiredKeyedService<IChatClient>(chatClientServiceKey);
-            return new ChatClientAgent(chatClient, instructions, key);
+            var tools = GetRegisteredToolsForAgent(sp, name);
+            return new ChatClientAgent(chatClient, instructions, key, tools: tools);
         });
     }
 
@@ -86,7 +93,8 @@ public static class AgentHostingServiceCollectionExtensions
         return services.AddAIAgent(name, (sp, key) =>
         {
             var chatClient = chatClientServiceKey is null ? sp.GetRequiredService<IChatClient>() : sp.GetRequiredKeyedService<IChatClient>(chatClientServiceKey);
-            return new ChatClientAgent(chatClient, instructions: instructions, name: key, description: description);
+            var tools = GetRegisteredToolsForAgent(sp, name);
+            return new ChatClientAgent(chatClient, instructions: instructions, name: key, description: description, tools: tools);
         });
     }
 
@@ -142,4 +150,10 @@ public static class AgentHostingServiceCollectionExtensions
         services.Add(ServiceDescriptor.Singleton(agentHostBuilderContext));
         services.AddSingleton<AgentCatalog, LocalAgentCatalog>();
     }
+
+    private static IList<AITool> GetRegisteredToolsForAgent(IServiceProvider serviceProvider, string agentName)
+    {
+        var registry = serviceProvider.GetService<LocalAgentToolRegistry>();
+        return registry?.GetTools(agentName) ?? [];
+    }
 }

dotnet/src/Microsoft.Agents.AI.Hosting/HostedAgentBuilderExtensions.cs

@@ -1,6 +1,9 @@
 // Copyright (c) Microsoft. All rights reserved.
 
 using System;
+using System.Linq;
+using Microsoft.Agents.AI.Hosting.Local;
+using Microsoft.Extensions.AI;
 using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Shared.Diagnostics;
 
@@ -59,4 +62,52 @@ public static class HostedAgentBuilderExtensions
         });
         return builder;
     }
+
+    /// <summary>
+    /// Adds an AI tool to an agent being configured with the service collection.
+    /// </summary>
+    /// <param name="builder">The hosted agent builder.</param>
+    /// <param name="tool">The AI tool to add to the agent.</param>
+    /// <returns>The same <see cref="IHostedAgentBuilder"/> instance so that additional calls can be chained.</returns>
+    /// <exception cref="ArgumentNullException">Thrown when <paramref name="builder"/> or <paramref name="tool"/> is <see langword="null"/>.</exception>
+    public static IHostedAgentBuilder WithAITool(this IHostedAgentBuilder builder, AITool tool)
+    {
+        Throw.IfNull(builder);
+        Throw.IfNull(tool);
+
+        var agentName = builder.Name;
+        var services = builder.ServiceCollection;
+
+        // Get or create the agent tool registry
+        var descriptor = services.FirstOrDefault(sd => !sd.IsKeyedService && sd.ServiceType.Equals(typeof(LocalAgentToolRegistry)));
+        if (descriptor?.ImplementationInstance is not LocalAgentToolRegistry toolRegistry)
+        {
+            toolRegistry = new();
+            services.Add(ServiceDescriptor.Singleton(toolRegistry));
+        }
+
+        toolRegistry.AddTool(agentName, tool);
+
+        return builder;
+    }
+
+    /// <summary>
+    /// Adds multiple AI tools to an agent being configured with the service collection.
+    /// </summary>
+    /// <param name="builder">The hosted agent builder.</param>
+    /// <param name="tools">The collection of AI tools to add to the agent.</param>
+    /// <returns>The same <see cref="IHostedAgentBuilder"/> instance so that additional calls can be chained.</returns>
+    /// <exception cref="ArgumentNullException">Thrown when <paramref name="builder"/> or <paramref name="tools"/> is <see langword="null"/>.</exception>
+    public static IHostedAgentBuilder WithAITools(this IHostedAgentBuilder builder, params AITool[] tools)
+    {
+        Throw.IfNull(builder);
+        Throw.IfNull(tools);
+
+        foreach (var tool in tools)
+        {
+            builder.WithAITool(tool);
+        }
+
+        return builder;
+    }
 }

dotnet/src/Microsoft.Agents.AI.Hosting/Local/LocalAgentToolRegistry.cs

@@ -0,0 +1,27 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System.Collections.Generic;
+using Microsoft.Extensions.AI;
+
+namespace Microsoft.Agents.AI.Hosting.Local;
+
+internal sealed class LocalAgentToolRegistry
+{
+    private readonly Dictionary<string, List<AITool>> _toolsByAgentName = new();
+
+    public void AddTool(string agentName, AITool tool)
+    {
+        if (!this._toolsByAgentName.TryGetValue(agentName, out var tools))
+        {
+            tools = [];
+            this._toolsByAgentName[agentName] = tools;
+        }
+
+        tools.Add(tool);
+    }
+
+    public IList<AITool> GetTools(string agentName)
+    {
+        return this._toolsByAgentName.TryGetValue(agentName, out var tools) ? tools : [];
+    }
+}

dotnet/src/Microsoft.Agents.AI/Microsoft.Agents.AI.csproj

@@ -31,8 +31,10 @@
   </PropertyGroup>
 
   <ItemGroup>
-    <InternalsVisibleTo Include="Microsoft.Agents.AI.UnitTests" />
     <InternalsVisibleTo Include="DynamicProxyGenAssembly2" />
+
+    <InternalsVisibleTo Include="Microsoft.Agents.AI.UnitTests" />
+    <InternalsVisibleTo Include="Microsoft.Agents.AI.Hosting.UnitTests"/>
   </ItemGroup>
 
 </Project>

dotnet/tests/Microsoft.Agents.AI.Hosting.UnitTests/HostedAgentBuilderToolsExtensionsTests.cs

@@ -0,0 +1,150 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System;
+using System.Collections.Generic;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.Extensions.AI;
+using Microsoft.Extensions.DependencyInjection;
+
+namespace Microsoft.Agents.AI.Hosting.UnitTests;
+
+/// <summary>
+/// Unit tests for AI tool registration extensions on <see cref="IHostedAgentBuilder"/>.
+/// </summary>
+public sealed class HostedAgentBuilderToolsExtensionsTests
+{
+    [Fact]
+    public void WithAITool_ThrowsWhenBuilderIsNull()
+    {
+        // Arrange
+        var tool = new DummyAITool();
+
+        // Act & Assert
+        Assert.Throws<ArgumentNullException>(() => HostedAgentBuilderExtensions.WithAITool(null!, tool));
+    }
+
+    [Fact]
+    public void WithAITool_ThrowsWhenToolIsNull()
+    {
+        // Arrange
+        var services = new ServiceCollection();
+        var builder = services.AddAIAgent("test-agent", "Test instructions");
+
+        // Act & Assert
+        Assert.Throws<ArgumentNullException>(() => builder.WithAITool(null!));
+    }
+
+    [Fact]
+    public void WithAITools_ThrowsWhenBuilderIsNull()
+    {
+        // Arrange
+        var tools = new[] { new DummyAITool() };
+
+        // Act & Assert
+        Assert.Throws<ArgumentNullException>(() => HostedAgentBuilderExtensions.WithAITools(null!, tools));
+    }
+
+    [Fact]
+    public void WithAITools_ThrowsWhenToolsArrayIsNull()
+    {
+        // Arrange
+        var services = new ServiceCollection();
+        var builder = services.AddAIAgent("test-agent", "Test instructions");
+
+        // Act & Assert
+        Assert.Throws<ArgumentNullException>(() => builder.WithAITools(null!));
+    }
+
+    [Fact]
+    public void RegisteredTools_ResolvesAllToolsForAgent()
+    {
+        // Arrange
+        var services = new ServiceCollection();
+        services.AddSingleton<IChatClient>(new MockChatClient());
+
+        var builder = services.AddAIAgent("test-agent", "Test instructions");
+        var tool1 = new DummyAITool();
+        var tool2 = new DummyAITool();
+
+        builder
+            .WithAITool(tool1)
+            .WithAITool(tool2);
+
+        var serviceProvider = services.BuildServiceProvider();
+
+        var agent1Tools = ResolveAgentTools(serviceProvider, "test-agent");
+        Assert.Contains(tool1, agent1Tools);
+        Assert.Contains(tool2, agent1Tools);
+    }
+
+    [Fact]
+    public void RegisteredTools_IsolatedPerAgent()
+    {
+        var services = new ServiceCollection();
+        services.AddSingleton<IChatClient>(new MockChatClient());
+
+        var builder1 = services.AddAIAgent("agent1", "Agent 1 instructions");
+        var builder2 = services.AddAIAgent("agent2", "Agent 2 instructions");
+
+        var tool1 = new DummyAITool();
+        var tool2 = new DummyAITool();
+        var tool3 = new DummyAITool();
+
+        builder1
+            .WithAITool(tool1)
+            .WithAITool(tool2);
+
+        builder2
+            .WithAITool(tool3);
+
+        var serviceProvider = services.BuildServiceProvider();
+
+        var agent1Tools = ResolveAgentTools(serviceProvider, "agent1");
+        var agent2Tools = ResolveAgentTools(serviceProvider, "agent2");
+
+        Assert.Contains(tool1, agent1Tools);
+        Assert.Contains(tool2, agent1Tools);
+        Assert.Contains(tool3, agent2Tools);
+    }
+
+    private static IList<AITool> ResolveAgentTools(IServiceProvider serviceProvider, string name)
+    {
+        var agent = serviceProvider.GetRequiredKeyedService<AIAgent>(name) as ChatClientAgent;
+        Assert.NotNull(agent?.ChatOptions?.Tools);
+        return agent.ChatOptions.Tools;
+    }
+
+    /// <summary>
+    /// Dummy AITool implementation for testing.
+    /// </summary>
+    private sealed class DummyAITool : AITool
+    {
+    }
+
+    /// <summary>
+    /// Mock chat client for testing.
+    /// </summary>
+    private sealed class MockChatClient : IChatClient
+    {
+        public Task<ChatResponse> GetResponseAsync(IEnumerable<ChatMessage> messages, ChatOptions? options = null, CancellationToken cancellationToken = default)
+        {
+            throw new NotImplementedException();
+        }
+
+        public IAsyncEnumerable<ChatResponseUpdate> GetStreamingResponseAsync(IEnumerable<ChatMessage> messages, ChatOptions? options = null, CancellationToken cancellationToken = default)
+        {
+            throw new NotImplementedException();
+        }
+
+        public object? GetService(Type serviceType, object? serviceKey = null)
+        {
+            return null;
+        }
+
+        public void Dispose()
+        {
+            throw new NotImplementedException();
+        }
+    }
+}