Add generic keyed version

src/Libraries/Microsoft.Extensions.AI/ChatCompletion/ChatClientBuilderServiceCollectionExtensions.cs

@@ -59,6 +59,18 @@ public static ChatClientBuilder AddKeyedChatClient(
         IChatClient innerClient)
         => AddKeyedChatClient(serviceCollection, serviceKey, _ => innerClient);
 
+    /// <summary>Adds a chat client to the <see cref="IServiceCollection"/>.</summary>
+    /// <typeparam name="T">The type of the inner <see cref="IChatClient"/> that represents the underlying backend. This will be resolved from the service provider.</typeparam>
+    /// <param name="serviceCollection">The <see cref="IServiceCollection"/> to which the client should be added.</param>
+    /// <param name="serviceKey">The key with which to associate the client.</param>
+    /// <returns>A <see cref="ChatClientBuilder"/> that can be used to build a pipeline around the inner client.</returns>
+    /// <remarks>The client is registered as a scoped service.</remarks>
+    public static ChatClientBuilder AddKeyedChatClient<T>(
+        this IServiceCollection serviceCollection,
+        object serviceKey)
+        where T : IChatClient
+        => AddKeyedChatClient(serviceCollection, serviceKey, services => services.GetRequiredService<T>());
+
     /// <summary>Adds a chat client to the <see cref="IServiceCollection"/>.</summary>
     /// <param name="serviceCollection">The <see cref="IServiceCollection"/> to which the client should be added.</param>
     /// <param name="serviceKey">The key with which to associate the client.</param>
@@ -75,7 +87,7 @@ public static ChatClientBuilder AddKeyedChatClient(
         _ = Throw.IfNull(innerClientFactory);
 
         var builder = new ChatClientBuilder(innerClientFactory);
-        _ = serviceCollection.AddKeyedSingleton(serviceKey, builder.Build);
+        _ = serviceCollection.AddKeyedSingleton(serviceKey, (services, _) => builder.Build(services));
         return builder;
     }
 }

test/Libraries/Microsoft.Extensions.AI.Tests/ChatCompletion/DependencyInjectionPatterns.cs

@@ -82,6 +82,83 @@ public void CanRegisterSingletonUsingSharedInstance()
         Assert.IsType<TestChatClient>(instance.InnerClient);
     }
 
+    [Fact]
+    public void CanRegisterKeyedSingletonUsingGenericType()
+    {
+        // Arrange/Act
+        ServiceCollection.AddSingleton(services => new TestChatClient { Services = services });
+        ServiceCollection.AddKeyedChatClient<TestChatClient>("mykey")
+            .UseSingletonMiddleware();
+
+        // Assert
+        var services = ServiceCollection.BuildServiceProvider();
+        using var scope1 = services.CreateScope();
+        using var scope2 = services.CreateScope();
+
+        Assert.Null(services.GetService<IChatClient>());
+
+        var instance1 = scope1.ServiceProvider.GetRequiredKeyedService<IChatClient>("mykey");
+        var instance1Copy = scope1.ServiceProvider.GetRequiredKeyedService<IChatClient>("mykey");
+        var instance2 = scope2.ServiceProvider.GetRequiredKeyedService<IChatClient>("mykey");
+
+        // Each scope gets the same instance, because it's singleton
+        var instance = Assert.IsType<SingletonMiddleware>(instance1);
+        Assert.Same(instance, instance1Copy);
+        Assert.Same(instance, instance2);
+        Assert.IsType<TestChatClient>(instance.InnerClient);
+    }
+
+    [Fact]
+    public void CanRegisterKeyedSingletonUsingFactory()
+    {
+        // Arrange/Act
+        ServiceCollection.AddKeyedChatClient("mykey", services => new TestChatClient { Services = services })
+            .UseSingletonMiddleware();
+
+        // Assert
+        var services = ServiceCollection.BuildServiceProvider();
+        using var scope1 = services.CreateScope();
+        using var scope2 = services.CreateScope();
+
+        Assert.Null(services.GetService<IChatClient>());
+
+        var instance1 = scope1.ServiceProvider.GetRequiredKeyedService<IChatClient>("mykey");
+        var instance1Copy = scope1.ServiceProvider.GetRequiredKeyedService<IChatClient>("mykey");
+        var instance2 = scope2.ServiceProvider.GetRequiredKeyedService<IChatClient>("mykey");
+
+        // Each scope gets the same instance, because it's singleton
+        var instance = Assert.IsType<SingletonMiddleware>(instance1);
+        Assert.Same(instance, instance1Copy);
+        Assert.Same(instance, instance2);
+        Assert.IsType<TestChatClient>(instance.InnerClient);
+    }
+
+    [Fact]
+    public void CanRegisterKeyedSingletonUsingSharedInstance()
+    {
+        // Arrange/Act
+        using var singleton = new TestChatClient();
+        ServiceCollection.AddKeyedChatClient("mykey", singleton)
+            .UseSingletonMiddleware();
+
+        // Assert
+        var services = ServiceCollection.BuildServiceProvider();
+        using var scope1 = services.CreateScope();
+        using var scope2 = services.CreateScope();
+
+        Assert.Null(services.GetService<IChatClient>());
+
+        var instance1 = scope1.ServiceProvider.GetRequiredKeyedService<IChatClient>("mykey");
+        var instance1Copy = scope1.ServiceProvider.GetRequiredKeyedService<IChatClient>("mykey");
+        var instance2 = scope2.ServiceProvider.GetRequiredKeyedService<IChatClient>("mykey");
+
+        // Each scope gets the same instance, because it's singleton
+        var instance = Assert.IsType<SingletonMiddleware>(instance1);
+        Assert.Same(instance, instance1Copy);
+        Assert.Same(instance, instance2);
+        Assert.IsType<TestChatClient>(instance.InnerClient);
+    }
+
     public class SingletonMiddleware(IServiceProvider services, IChatClient inner) : DelegatingChatClient(inner)
     {
         public new IChatClient InnerClient => base.InnerClient;