Add extensible chat client routing

src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingChatClient.cs

@@ -0,0 +1,127 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System;
+using System.Collections.Generic;
+using System.Diagnostics.CodeAnalysis;
+using System.Runtime.CompilerServices;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.Shared.DiagnosticIds;
+using Microsoft.Shared.Diagnostics;
+
+namespace Microsoft.Extensions.AI;
+
+/// <summary>
+/// Provides a template for an <see cref="IChatClient"/> that selects and invokes another chat client.
+/// </summary>
+/// <remarks>
+/// Derived classes implement <see cref="SelectClientAsync"/> to supply one client for each request. The selected
+/// client is invoked once, and its response or failure is propagated to the caller.
+/// The exact client instance represents the selected routing identity. Caller-supplied options may vary per
+/// invocation, but they are ephemeral and do not participate in that identity. Applications can use distinct
+/// configured client wrappers when configurations require distinct routing identities, while custom policies may
+/// maintain separate application-specific grouping keys.
+/// </remarks>
+[Experimental(DiagnosticIds.Experiments.AIRoutingChat, UrlFormat = DiagnosticIds.UrlFormat)]
+public abstract class RoutingChatClient : IChatClient
+{
+    /// <summary>Creates a routing client that selects one client for each request.</summary>
+    /// <param name="clientSelector">The callback that selects the client to invoke.</param>
+    /// <returns>A routing client that uses <paramref name="clientSelector"/>.</returns>
+    /// <remarks>The selected clients are caller-owned and are not disposed by the returned routing client.</remarks>
+    /// <exception cref="ArgumentNullException"><paramref name="clientSelector"/> is <see langword="null"/>.</exception>
+    public static RoutingChatClient Create(
+        Func<RoutingContext, CancellationToken, ValueTask<IChatClient>> clientSelector)
+    {
+        _ = Throw.IfNull(clientSelector);
+
+        return new CallbackRoutingChatClient(clientSelector);
+    }
+
+    /// <summary>Selects the client to invoke for the request.</summary>
+    /// <param name="context">The request-specific inputs.</param>
+    /// <param name="cancellationToken">The cancellation token supplied for the request.</param>
+    /// <returns>The client to invoke.</returns>
+    /// <remarks>
+    /// Client-specific behavior should generally be attached to the returned client. Exceptions from this method
+    /// propagate to the caller.
+    /// </remarks>
+    protected abstract ValueTask<IChatClient> SelectClientAsync(
+        RoutingContext context,
+        CancellationToken cancellationToken);
+
+    /// <inheritdoc/>
+    public virtual async Task<ChatResponse> GetResponseAsync(
+        IEnumerable<ChatMessage> messages, ChatOptions? options = null, CancellationToken cancellationToken = default)
+    {
+        _ = Throw.IfNull(messages);
+
+        var context = new RoutingContext(messages, options);
+        IChatClient client = await SelectClientAsync(context, cancellationToken).ConfigureAwait(false) ??
+            throw new InvalidOperationException($"{nameof(SelectClientAsync)} returned null.");
+
+        return await client.GetResponseAsync(
+            context.Messages,
+            options?.Clone(),
+            cancellationToken).ConfigureAwait(false);
+    }
+
+    /// <inheritdoc/>
+    public virtual async IAsyncEnumerable<ChatResponseUpdate> GetStreamingResponseAsync(
+        IEnumerable<ChatMessage> messages, ChatOptions? options = null,
+        [EnumeratorCancellation] CancellationToken cancellationToken = default)
+    {
+        _ = Throw.IfNull(messages);
+
+        var context = new RoutingContext(messages, options);
+        IChatClient client = await SelectClientAsync(context, cancellationToken).ConfigureAwait(false) ??
+            throw new InvalidOperationException($"{nameof(SelectClientAsync)} returned null.");
+
+        await foreach (ChatResponseUpdate update in
+            client.GetStreamingResponseAsync(context.Messages, options?.Clone(), cancellationToken)
+                .WithCancellation(cancellationToken)
+                .ConfigureAwait(false))
+        {
+            yield return update;
+        }
+    }
+
+    /// <inheritdoc/>
+    public virtual object? GetService(Type serviceType, object? serviceKey = null)
+    {
+        _ = Throw.IfNull(serviceType);
+
+        return serviceKey is null && serviceType.IsInstanceOfType(this) ? this : null;
+    }
+
+    /// <inheritdoc/>
+    public void Dispose()
+    {
+        Dispose(disposing: true);
+        GC.SuppressFinalize(this);
+    }
+
+    /// <summary>Provides a mechanism for releasing resources owned by the derived instance.</summary>
+    /// <param name="disposing"><see langword="true"/> when called from <see cref="Dispose()"/>.</param>
+    /// <remarks>The default implementation performs no operation.</remarks>
+    protected virtual void Dispose(bool disposing)
+    {
+    }
+
+    private sealed class CallbackRoutingChatClient : RoutingChatClient
+    {
+        private readonly Func<RoutingContext, CancellationToken, ValueTask<IChatClient>> _clientSelector;
+
+        public CallbackRoutingChatClient(
+            Func<RoutingContext, CancellationToken, ValueTask<IChatClient>> clientSelector)
+        {
+            _clientSelector = clientSelector;
+        }
+
+        protected override ValueTask<IChatClient> SelectClientAsync(
+            RoutingContext context,
+            CancellationToken cancellationToken) =>
+            _clientSelector(context, cancellationToken);
+    }
+}

src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingContext.cs

@@ -0,0 +1,50 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System.Collections.Generic;
+using System.Diagnostics.CodeAnalysis;
+using Microsoft.Shared.DiagnosticIds;
+using Microsoft.Shared.Diagnostics;
+
+namespace Microsoft.Extensions.AI;
+
+/// <summary>Provides request-specific inputs to a <see cref="RoutingChatClient"/>.</summary>
+/// <remarks>
+/// <para>
+/// One context is created for each call to <see cref="IChatClient.GetResponseAsync"/> and for each enumeration
+/// started from the sequence returned by <see cref="IChatClient.GetStreamingResponseAsync"/>.
+/// </para>
+/// <para>
+/// <see cref="ChatOptions"/> is cloned when the context is created. It is provided for client selection and is
+/// independent of both the caller's instance and the options passed to the selected client.
+/// </para>
+/// </remarks>
+[Experimental(DiagnosticIds.Experiments.AIRoutingChat, UrlFormat = DiagnosticIds.UrlFormat)]
+public class RoutingContext
+{
+    /// <summary>Initializes a new instance of the <see cref="RoutingContext"/> class.</summary>
+    /// <param name="messages">The messages to route.</param>
+    /// <param name="chatOptions">The options associated with this context.</param>
+    public RoutingContext(
+        IEnumerable<ChatMessage> messages,
+        ChatOptions? chatOptions)
+    {
+        _ = Throw.IfNull(messages);
+
+        Messages = messages;
+        ChatOptions = chatOptions?.Clone();
+    }
+
+    /// <summary>Gets the messages supplied to client selection and the selected client.</summary>
+    /// <remarks>
+    /// Selection and failover may enumerate this sequence multiple times. Callers should supply a repeatable sequence.
+    /// </remarks>
+    public IEnumerable<ChatMessage> Messages { get; }
+
+    /// <summary>Gets a snapshot of the request options supplied to client selection.</summary>
+    /// <remarks>
+    /// Changes do not affect the caller's instance or the options passed to the selected client. Client-specific
+    /// behavior should generally be attached to the returned client.
+    /// </remarks>
+    public ChatOptions? ChatOptions { get; }
+}

src/Libraries/Microsoft.Extensions.AI.Abstractions/Microsoft.Extensions.AI.Abstractions.json

@@ -3966,6 +3966,64 @@
         }
       ]
     },
+    {
+      "Type": "abstract class Microsoft.Extensions.AI.RoutingChatClient : Microsoft.Extensions.AI.IChatClient, System.IDisposable",
+      "Stage": "Experimental",
+      "Methods": [
+        {
+          "Member": "Microsoft.Extensions.AI.RoutingChatClient.RoutingChatClient();",
+          "Stage": "Experimental"
+        },
+        {
+          "Member": "static Microsoft.Extensions.AI.RoutingChatClient Microsoft.Extensions.AI.RoutingChatClient.Create(System.Func<Microsoft.Extensions.AI.RoutingContext, System.Threading.CancellationToken, System.Threading.Tasks.ValueTask<Microsoft.Extensions.AI.IChatClient>> clientSelector);",
+          "Stage": "Experimental"
+        },
+        {
+          "Member": "void Microsoft.Extensions.AI.RoutingChatClient.Dispose();",
+          "Stage": "Experimental"
+        },
+        {
+          "Member": "virtual void Microsoft.Extensions.AI.RoutingChatClient.Dispose(bool disposing);",
+          "Stage": "Experimental"
+        },
+        {
+          "Member": "virtual System.Threading.Tasks.Task<Microsoft.Extensions.AI.ChatResponse> Microsoft.Extensions.AI.RoutingChatClient.GetResponseAsync(System.Collections.Generic.IEnumerable<Microsoft.Extensions.AI.ChatMessage> messages, Microsoft.Extensions.AI.ChatOptions? options = null, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken));",
+          "Stage": "Experimental"
+        },
+        {
+          "Member": "virtual object? Microsoft.Extensions.AI.RoutingChatClient.GetService(System.Type serviceType, object? serviceKey = null);",
+          "Stage": "Experimental"
+        },
+        {
+          "Member": "virtual System.Collections.Generic.IAsyncEnumerable<Microsoft.Extensions.AI.ChatResponseUpdate> Microsoft.Extensions.AI.RoutingChatClient.GetStreamingResponseAsync(System.Collections.Generic.IEnumerable<Microsoft.Extensions.AI.ChatMessage> messages, Microsoft.Extensions.AI.ChatOptions? options = null, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken));",
+          "Stage": "Experimental"
+        },
+        {
+          "Member": "abstract System.Threading.Tasks.ValueTask<Microsoft.Extensions.AI.IChatClient> Microsoft.Extensions.AI.RoutingChatClient.SelectClientAsync(Microsoft.Extensions.AI.RoutingContext context, System.Threading.CancellationToken cancellationToken);",
+          "Stage": "Experimental"
+        }
+      ]
+    },
+    {
+      "Type": "class Microsoft.Extensions.AI.RoutingContext",
+      "Stage": "Experimental",
+      "Methods": [
+        {
+          "Member": "Microsoft.Extensions.AI.RoutingContext.RoutingContext(System.Collections.Generic.IEnumerable<Microsoft.Extensions.AI.ChatMessage> messages, Microsoft.Extensions.AI.ChatOptions? chatOptions);",
+          "Stage": "Experimental"
+        }
+      ],
+      "Properties": [
+        {
+          "Member": "Microsoft.Extensions.AI.ChatOptions? Microsoft.Extensions.AI.RoutingContext.ChatOptions { get; }",
+          "Stage": "Experimental"
+        },
+        {
+          "Member": "System.Collections.Generic.IEnumerable<Microsoft.Extensions.AI.ChatMessage> Microsoft.Extensions.AI.RoutingContext.Messages { get; }",
+          "Stage": "Experimental"
+        }
+      ]
+    },
     {
       "Type": "class Microsoft.Extensions.AI.SessionUpdateRealtimeClientMessage : Microsoft.Extensions.AI.RealtimeClientMessage",
       "Stage": "Experimental",