diff --git a/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingChatClient.cs b/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingChatClient.cs new file mode 100644 index 00000000000..0f28c518c3b --- /dev/null +++ b/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; + +/// +/// Provides a template for an that selects and invokes another chat client. +/// +/// +/// Derived classes implement to supply one selection for each request. The selected +/// client is invoked once with the selected options, and its response or failure is propagated to the caller. +/// +[Experimental(DiagnosticIds.Experiments.AIRoutingChat, UrlFormat = DiagnosticIds.UrlFormat)] +public abstract class RoutingChatClient : IChatClient +{ + /// Creates a routing client that selects one client and its options for each request. + /// The callback that produces the selection to invoke. + /// A routing client that uses . + /// The selected clients are caller-owned and are not disposed by the returned routing client. + /// is . + public static RoutingChatClient Create( + Func> selector) + { + _ = Throw.IfNull(selector); + + return new CallbackRoutingChatClient(selector); + } + + /// Selects the client and complete options to use for the request. + /// The request-specific inputs. + /// The cancellation token supplied for the request. + /// The selection to invoke. + /// + /// The returned are complete attempt options, not an overlay. + /// Implementations may clone and adjust the clone for the selected + /// attempt. Exceptions from this method propagate to the caller. + /// + protected abstract ValueTask SelectClientAsync( + RoutingContext context, + CancellationToken cancellationToken); + + /// + public virtual async Task GetResponseAsync( + IEnumerable messages, ChatOptions? options = null, CancellationToken cancellationToken = default) + { + _ = Throw.IfNull(messages); + + var context = new RoutingContext(messages, options); + RoutingSelection selection = await SelectClientAsync(context, cancellationToken).ConfigureAwait(false) ?? + throw new InvalidOperationException($"{nameof(SelectClientAsync)} returned null."); + + return await selection.Client.GetResponseAsync( + context.Messages, + selection.ChatOptions, + cancellationToken).ConfigureAwait(false); + } + + /// + public virtual async IAsyncEnumerable GetStreamingResponseAsync( + IEnumerable messages, ChatOptions? options = null, + [EnumeratorCancellation] CancellationToken cancellationToken = default) + { + _ = Throw.IfNull(messages); + + var context = new RoutingContext(messages, options); + RoutingSelection selection = await SelectClientAsync(context, cancellationToken).ConfigureAwait(false) ?? + throw new InvalidOperationException($"{nameof(SelectClientAsync)} returned null."); + + await foreach (ChatResponseUpdate update in + selection.Client.GetStreamingResponseAsync( + context.Messages, + selection.ChatOptions, + cancellationToken) + .WithCancellation(cancellationToken) + .ConfigureAwait(false)) + { + yield return update; + } + } + + /// + public virtual object? GetService(Type serviceType, object? serviceKey = null) + { + _ = Throw.IfNull(serviceType); + + return serviceKey is null && serviceType.IsInstanceOfType(this) ? this : null; + } + + /// + public void Dispose() + { + Dispose(disposing: true); + GC.SuppressFinalize(this); + } + + /// Provides a mechanism for releasing resources owned by the derived instance. + /// when called from . + /// The default implementation performs no operation. + protected virtual void Dispose(bool disposing) + { + } + + private sealed class CallbackRoutingChatClient : RoutingChatClient + { + private readonly Func> _selector; + + public CallbackRoutingChatClient( + Func> selector) + { + _selector = selector; + } + + protected override ValueTask SelectClientAsync( + RoutingContext context, + CancellationToken cancellationToken) => + _selector(context, cancellationToken); + } +} diff --git a/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingContext.cs b/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingContext.cs new file mode 100644 index 00000000000..1c72a4d18fb --- /dev/null +++ b/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingContext.cs @@ -0,0 +1,49 @@ +// 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; + +/// Provides request-specific inputs to a . +/// +/// +/// One context is created for each call to and for each enumeration +/// started from the sequence returned by . +/// +/// +/// is cloned when the context is created and represents the request baseline. Selectors can +/// clone and modify it to produce complete attempt options through . +/// +/// +[Experimental(DiagnosticIds.Experiments.AIRoutingChat, UrlFormat = DiagnosticIds.UrlFormat)] +public class RoutingContext +{ + /// Initializes a new instance of the class. + /// The messages to route. + /// The options associated with this context. + public RoutingContext( + IEnumerable messages, + ChatOptions? chatOptions) + { + _ = Throw.IfNull(messages); + + Messages = messages; + ChatOptions = chatOptions?.Clone(); + } + + /// Gets the messages supplied to client selection and the selected client. + /// + /// Selection and failover may enumerate this sequence multiple times. Callers should supply a repeatable sequence. + /// + public IEnumerable Messages { get; } + + /// Gets the request options supplied to client selection. + /// + /// This is a clone of the caller's options. Selectors should clone it before applying attempt-specific changes. + /// + public ChatOptions? ChatOptions { get; } +} diff --git a/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingSelection.cs b/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingSelection.cs new file mode 100644 index 00000000000..5f8947b83cd --- /dev/null +++ b/src/Libraries/Microsoft.Extensions.AI.Abstractions/ChatRouting/RoutingSelection.cs @@ -0,0 +1,35 @@ +// 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.Diagnostics.CodeAnalysis; +using Microsoft.Shared.DiagnosticIds; +using Microsoft.Shared.Diagnostics; + +namespace Microsoft.Extensions.AI; + +/// Represents the client and complete options selected for one routing attempt. +/// +/// The exact instance represents the routing identity. are ephemeral +/// attempt data and do not participate in that identity. The selected client may further transform the options. +/// +[Experimental(DiagnosticIds.Experiments.AIRoutingChat, UrlFormat = DiagnosticIds.UrlFormat)] +public sealed class RoutingSelection +{ + /// Initializes a new instance of the class. + /// The client selected for the attempt. + /// The complete options selected for the attempt, or . + /// is . + public RoutingSelection(IChatClient client, ChatOptions? chatOptions) + { + Client = Throw.IfNull(client); + ChatOptions = chatOptions?.Clone(); + } + + /// Gets the client selected for the attempt. + public IChatClient Client { get; } + + /// Gets the complete options selected for the attempt. + /// This is a clone of the options supplied to the constructor. + public ChatOptions? ChatOptions { get; } +} diff --git a/src/Libraries/Microsoft.Extensions.AI.Abstractions/Microsoft.Extensions.AI.Abstractions.json b/src/Libraries/Microsoft.Extensions.AI.Abstractions/Microsoft.Extensions.AI.Abstractions.json index dc4542d4ae7..0518ee08b60 100644 --- a/src/Libraries/Microsoft.Extensions.AI.Abstractions/Microsoft.Extensions.AI.Abstractions.json +++ b/src/Libraries/Microsoft.Extensions.AI.Abstractions/Microsoft.Extensions.AI.Abstractions.json @@ -3966,6 +3966,84 @@ } ] }, + { + "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> selector);", + "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.RoutingChatClient.GetResponseAsync(System.Collections.Generic.IEnumerable 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.RoutingChatClient.GetStreamingResponseAsync(System.Collections.Generic.IEnumerable 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.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 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.RoutingContext.Messages { get; }", + "Stage": "Experimental" + } + ] + }, + { + "Type": "sealed class Microsoft.Extensions.AI.RoutingSelection", + "Stage": "Experimental", + "Methods": [ + { + "Member": "Microsoft.Extensions.AI.RoutingSelection.RoutingSelection(Microsoft.Extensions.AI.IChatClient client, Microsoft.Extensions.AI.ChatOptions? chatOptions);", + "Stage": "Experimental" + } + ], + "Properties": [ + { + "Member": "Microsoft.Extensions.AI.ChatOptions? Microsoft.Extensions.AI.RoutingSelection.ChatOptions { get; }", + "Stage": "Experimental" + }, + { + "Member": "Microsoft.Extensions.AI.IChatClient Microsoft.Extensions.AI.RoutingSelection.Client { get; }", + "Stage": "Experimental" + } + ] + }, { "Type": "class Microsoft.Extensions.AI.SessionUpdateRealtimeClientMessage : Microsoft.Extensions.AI.RealtimeClientMessage", "Stage": "Experimental", diff --git a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClient.cs b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClient.cs new file mode 100644 index 00000000000..b85268e13f0 --- /dev/null +++ b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClient.cs @@ -0,0 +1,376 @@ +// 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; +using System.Diagnostics.CodeAnalysis; +using System.Runtime.CompilerServices; +using System.Runtime.ExceptionServices; +using System.Threading; +using System.Threading.Tasks; +using Microsoft.Shared.DiagnosticIds; +using Microsoft.Shared.Diagnostics; + +namespace Microsoft.Extensions.AI; + +/// +/// Provides a template for a that can select another client after an invocation fails. +/// +/// +/// +/// The client for each attempt is supplied by . After an invocation, +/// reports the concrete attempt and whether another selection will follow. An +/// uncanceled failure causes another selection only when it happened before any streaming output was exposed and the +/// attempt limit permits it. +/// +/// +/// The base class owns invocation, streaming commitment, attempt limits, and attempt reporting. Derived classes own +/// client selection, policy state, selection-failure cleanup, and the lifetime of clients they retain. +/// +/// +/// Once streaming enumeration begins, callers must dispose the enumerator. Abandoning an active enumerator without +/// disposing it prevents both inner enumerator disposal and the terminal routing update. +/// +/// +[Experimental(DiagnosticIds.Experiments.AIRoutingChat, UrlFormat = DiagnosticIds.UrlFormat)] +public abstract class FailoverChatClient : RoutingChatClient +{ + /// Gets or sets the maximum number of client invocations permitted for one request. + /// + /// A positive attempt limit, or to leave termination to client selection and request + /// cancellation. The default is . + /// + /// + /// The value is captured when a non-streaming request begins or when a streaming response begins enumeration. + /// Changing it does not affect requests or enumerations already in progress. + /// + /// The value is not or positive. + public int? MaximumAttemptsPerRequest + { + get; + set + { + if (value is <= 0) + { + Throw.ArgumentOutOfRangeException(nameof(value)); + } + + field = value; + } + } + + /// Invoked after a client invocation completes, fails, or is abandoned. + /// The request-specific inputs. + /// The attempted client invocation. + /// + /// if the base will not select another client after this method returns successfully; + /// otherwise, . + /// + /// The cancellation token supplied for the request. + /// A task representing the update operation. + /// + /// + /// The default implementation performs no operation. A nonterminal update always contains an uncanceled, + /// pre-output failed attempt. State changes made by the override are visible to the next call to + /// . + /// + /// + /// This method is invoked once after each selected-client invocation, whether it completes, fails, or is abandoned. + /// Selection failures are not reported. A selector that retains request-scoped state must release it before + /// throwing; a request may therefore end without a terminal update when selection fails. + /// + /// + /// Exceptions from this method propagate to the caller. A terminal update exception replaces the response or + /// exception already produced by the request. A nonterminal update exception stops routing without another update. + /// An override that retains per-request state must release that state before throwing because no later update is + /// made after an update exception. + /// + /// + protected virtual ValueTask OnRoutingUpdateAsync( + RoutingContext context, + FailoverChatClientAttempt attempt, + bool isTerminal, + CancellationToken cancellationToken) => default; + + /// + public sealed override async Task GetResponseAsync( + IEnumerable messages, ChatOptions? options = null, CancellationToken cancellationToken = default) + { + _ = Throw.IfNull(messages); + + var context = new RoutingContext(messages, options); + int? maximumAttempts = MaximumAttemptsPerRequest; + int attemptCount = 0; + + while (true) + { + RoutingSelection selection = + await SelectClientAsync(context, cancellationToken).ConfigureAwait(false) ?? + throw new InvalidOperationException($"{nameof(SelectClientAsync)} returned null."); + IChatClient selectedClient = selection.Client; + ChatOptions? invocationOptions = selection.ChatOptions; + ChatOptions? attemptOptions = invocationOptions?.Clone(); + + attemptCount++; + ChatResponse? response = null; + Exception? exception = null; + Stopwatch stopwatch = Stopwatch.StartNew(); + + try + { + response = await selectedClient.GetResponseAsync( + context.Messages, + invocationOptions, + cancellationToken).ConfigureAwait(false); + } + catch (Exception ex) + { + exception = ex; + } + + stopwatch.Stop(); + var attempt = new FailoverChatClientAttempt( + selectedClient, + attemptOptions, + exception, + stopwatch.Elapsed, + timeToFirstUpdate: null, + responseCompleted: exception is null, + outputCommitted: false); + bool cancellationRequested = + exception is not null && + cancellationToken.IsCancellationRequested; + bool isTerminal = + exception is null || + cancellationRequested || + (maximumAttempts is int limit && attemptCount >= limit); + + await OnRoutingUpdateAsync(context, attempt, isTerminal, cancellationToken).ConfigureAwait(false); + + if (exception is null) + { + return response!; + } + + if (cancellationRequested) + { + cancellationToken.ThrowIfCancellationRequested(); + } + + if (isTerminal) + { + Rethrow(exception); + } + } + } + + /// + public sealed override async IAsyncEnumerable GetStreamingResponseAsync( + IEnumerable messages, ChatOptions? options = null, + [EnumeratorCancellation] CancellationToken cancellationToken = default) + { + _ = Throw.IfNull(messages); + + var context = new RoutingContext(messages, options); + int? maximumAttempts = MaximumAttemptsPerRequest; + int attemptCount = 0; + + while (true) + { + RoutingSelection selection = + await SelectClientAsync(context, cancellationToken).ConfigureAwait(false) ?? + throw new InvalidOperationException($"{nameof(SelectClientAsync)} returned null."); + IChatClient selectedClient = selection.Client; + ChatOptions? invocationOptions = selection.ChatOptions; + ChatOptions? attemptOptions = invocationOptions?.Clone(); + + attemptCount++; + bool reachedAttemptLimit = maximumAttempts is int limit && attemptCount >= limit; + IAsyncEnumerator? enumerator = null; + TimeSpan? timeToFirstUpdate = null; + bool hasCurrent; + + Stopwatch stopwatch = Stopwatch.StartNew(); + try + { + enumerator = selectedClient + .GetStreamingResponseAsync( + context.Messages, + invocationOptions, + cancellationToken) + .GetAsyncEnumerator(cancellationToken); + + hasCurrent = await enumerator.MoveNextAsync().ConfigureAwait(false); + } + catch (Exception ex) + { + stopwatch.Stop(); + Exception exception = (await DisposeEnumeratorAsync(enumerator, ex).ConfigureAwait(false))!; + var attempt = new FailoverChatClientAttempt( + selectedClient, + attemptOptions, + exception, + stopwatch.Elapsed, + timeToFirstUpdate: null, + responseCompleted: false, + outputCommitted: false); + bool cancellationRequested = cancellationToken.IsCancellationRequested; + bool isTerminal = cancellationRequested || reachedAttemptLimit; + + await OnRoutingUpdateAsync(context, attempt, isTerminal, cancellationToken).ConfigureAwait(false); + + if (cancellationRequested) + { + cancellationToken.ThrowIfCancellationRequested(); + } + + if (isTerminal) + { + Rethrow(exception); + } + + continue; + } + + stopwatch.Stop(); + bool responseCompleted = false; + bool outputCommitted = false; + bool isTerminalAttempt = false; + Exception? terminalException = null; + + try + { + while (hasCurrent) + { + stopwatch.Start(); + bool hasCurrentValue = TryGetCurrent( + enumerator, + out ChatResponseUpdate current, + out terminalException); + stopwatch.Stop(); + if (!hasCurrentValue) + { + break; + } + + timeToFirstUpdate ??= stopwatch.Elapsed; + outputCommitted = true; + yield return current; + + stopwatch.Start(); + try + { + hasCurrent = await enumerator.MoveNextAsync().ConfigureAwait(false); + } + catch (Exception ex) + { + terminalException = ex; + break; + } + finally + { + stopwatch.Stop(); + } + } + + responseCompleted = terminalException is null; + } + finally + { + terminalException = + await DisposeEnumeratorAsync(enumerator, terminalException).ConfigureAwait(false); + + var attempt = new FailoverChatClientAttempt( + selectedClient, + attemptOptions, + terminalException, + stopwatch.Elapsed, + timeToFirstUpdate, + responseCompleted: responseCompleted && terminalException is null, + outputCommitted: outputCommitted); + bool cancellationRequested = + terminalException is not null && + cancellationToken.IsCancellationRequested; + isTerminalAttempt = + attempt.ResponseCompleted || + outputCommitted || + cancellationRequested || + reachedAttemptLimit; + + await OnRoutingUpdateAsync( + context, + attempt, + isTerminalAttempt, + cancellationToken).ConfigureAwait(false); + + if (terminalException is not null) + { + if (cancellationRequested) + { + cancellationToken.ThrowIfCancellationRequested(); + } + + if (isTerminalAttempt) + { + Rethrow(terminalException); + } + } + } + + if (!isTerminalAttempt) + { + continue; + } + + yield break; + } + } + +#pragma warning disable EA0014 // IAsyncDisposable.DisposeAsync doesn't support cancellation. + private static async ValueTask DisposeEnumeratorAsync( + IAsyncDisposable? disposable, Exception? exception) +#pragma warning restore EA0014 + { + if (disposable is not null) + { + try + { + await disposable.DisposeAsync().ConfigureAwait(false); + } + catch (Exception ex) + { + return ex; + } + } + + return exception; + } + + private static bool TryGetCurrent( + IAsyncEnumerator enumerator, + out ChatResponseUpdate current, + out Exception? exception) + { + try + { + current = enumerator.Current; + exception = null; + return true; + } + catch (Exception ex) + { + current = null!; + exception = ex; + return false; + } + } + +#if NET + [DoesNotReturn] +#endif + private static void Rethrow(Exception exception) + { + ExceptionDispatchInfo.Capture(exception).Throw(); + } +} diff --git a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClientAttempt.cs b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClientAttempt.cs new file mode 100644 index 00000000000..21cb0e32585 --- /dev/null +++ b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/FailoverChatClientAttempt.cs @@ -0,0 +1,98 @@ +// 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.Diagnostics; +using System.Diagnostics.CodeAnalysis; +using Microsoft.Shared.DiagnosticIds; +using Microsoft.Shared.Diagnostics; + +namespace Microsoft.Extensions.AI; + +/// Represents one client invocation attempt performed by a . +/// +/// +/// A completed response has set to and +/// set to . +/// +/// +/// A failed invocation has set to and +/// set to the observed exception. If a streaming caller stops enumerating before the +/// response ends and disposal completes successfully, both and +/// are unset. +/// +/// +/// is independent of the outcome and indicates whether any streaming update was +/// exposed to the caller. +/// +/// +[Experimental(DiagnosticIds.Experiments.AIRoutingChat, UrlFormat = DiagnosticIds.UrlFormat)] +public sealed class FailoverChatClientAttempt +{ + internal FailoverChatClientAttempt( + IChatClient client, + ChatOptions? chatOptions, + Exception? exception, + TimeSpan duration, + TimeSpan? timeToFirstUpdate, + bool responseCompleted, + bool outputCommitted) + { + Debug.Assert(duration >= TimeSpan.Zero, "Expected a non-negative duration."); + Debug.Assert( + timeToFirstUpdate is not { } ttfu || (ttfu >= TimeSpan.Zero && ttfu <= duration), + "Expected time to first update to be within the active duration."); + Debug.Assert( + !responseCompleted || exception is null, + "A completed response should not have an exception."); + Debug.Assert( + outputCommitted == timeToFirstUpdate.HasValue, + "Output commitment and time to first update should agree."); + + Client = Throw.IfNull(client); + ChatOptions = chatOptions; + Exception = exception; + Duration = duration; + TimeToFirstUpdate = timeToFirstUpdate; + ResponseCompleted = responseCompleted; + OutputCommitted = outputCommitted; + } + + /// Gets the client that was invoked. + public IChatClient Client { get; } + + /// Gets a snapshot of the options passed by the router to the selected client. + /// + /// The selected client or middleware in its pipeline may subsequently clone or transform these options before + /// invoking an underlying provider. + /// + public ChatOptions? ChatOptions { get; } + + /// Gets the time spent actively invoking the client. + /// For streaming, time spent by the caller processing yielded updates is excluded. + public TimeSpan Duration { get; } + + /// Gets the exception observed while invoking or disposing the client, if any. + /// + /// If invocation and disposal both throw, this contains the disposal exception. + /// + /// A value does not necessarily indicate success; inspect + /// to distinguish a completed response from a streaming response that the caller stopped consuming. + /// + /// + public Exception? Exception { get; } + + /// Gets a value indicating whether any streaming update was exposed to the caller. + /// This is always for non-streaming invocations. + public bool OutputCommitted { get; } + + /// Gets a value indicating whether the response completed successfully. + /// + /// For streaming responses, this is when the caller stops enumeration before the + /// response stream ends. + /// + public bool ResponseCompleted { get; } + + /// Gets the time until the first streaming update, if this was a non-empty streaming invocation. + public TimeSpan? TimeToFirstUpdate { get; } +} diff --git a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/OrderedFailoverChatClient.cs b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/OrderedFailoverChatClient.cs new file mode 100644 index 00000000000..f90e3694d7a --- /dev/null +++ b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/OrderedFailoverChatClient.cs @@ -0,0 +1,173 @@ +// 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.Concurrent; +using System.Collections.Generic; +using System.Diagnostics.CodeAnalysis; +using System.Runtime.ExceptionServices; +using System.Threading; +using System.Threading.Tasks; +using Microsoft.Shared.DiagnosticIds; +using Microsoft.Shared.Diagnostics; + +namespace Microsoft.Extensions.AI; + +/// Provides ordered failover across a sequence of chat clients. +/// +/// +/// The clients are tried in order. An invocation failure before streaming output is exposed advances to the next +/// client. Cancellation and failures after streaming output is exposed are propagated without failover. +/// +/// +/// The configured clients are snapshotted by the constructor and must contain unique object references. When every +/// client has failed, the final failure is rethrown. +/// +/// +[Experimental(DiagnosticIds.Experiments.AIRoutingChat, UrlFormat = DiagnosticIds.UrlFormat)] +public sealed class OrderedFailoverChatClient : FailoverChatClient +{ + private readonly bool _leaveOpen; + private readonly IChatClient[] _clients; + private readonly ConcurrentDictionary _requestStates = new(); + private bool _disposed; + + /// Initializes a new instance of the class. + /// The clients to invoke, in fallback order. + /// + /// to leave inner clients open when this instance is disposed; + /// otherwise, . + /// + /// is . + /// + /// is empty, contains , or contains the same client instance more + /// than once. + /// + public OrderedFailoverChatClient(IReadOnlyList clients, bool leaveOpen = false) + { + _ = Throw.IfNull(clients); + + IChatClient[] clientsSnapshot = [.. clients]; + if (clientsSnapshot.Length == 0) + { + Throw.ArgumentException(nameof(clients), "At least one client must be provided."); + } + + for (int i = 0; i < clientsSnapshot.Length; i++) + { + if (clientsSnapshot[i] is null) + { + Throw.ArgumentException(nameof(clients), "Clients must not contain null."); + } + + for (int j = 0; j < i; j++) + { + if (ReferenceEquals(clientsSnapshot[j], clientsSnapshot[i])) + { + Throw.ArgumentException(nameof(clients), "Each client instance must be unique."); + } + } + } + + _leaveOpen = leaveOpen; + _clients = clientsSnapshot; + } + + /// + protected override ValueTask SelectClientAsync( + RoutingContext context, + CancellationToken cancellationToken) + { + _ = Throw.IfNull(context); + _ = cancellationToken; + + if (!_requestStates.TryRemove(context, out RequestState? state)) + { + return new(new RoutingSelection(_clients[0], context.ChatOptions)); + } + + if (state.ClientIndex < _clients.Length) + { + return new(new RoutingSelection(_clients[state.ClientIndex], context.ChatOptions)); + } + + ExceptionDispatchInfo.Capture(state.LastException).Throw(); + throw state.LastException; + } + + /// + protected override ValueTask OnRoutingUpdateAsync( + RoutingContext context, + FailoverChatClientAttempt attempt, + bool isTerminal, + CancellationToken cancellationToken) + { + _ = cancellationToken; + + if (isTerminal) + { + _ = _requestStates.TryRemove(context, out _); + return default; + } + + if (attempt.Exception is null) + { + _ = _requestStates.TryRemove(context, out _); + throw new InvalidOperationException("A nonterminal routing update requires a failed client invocation."); + } + + int clientIndex = IndexOfClient(attempt.Client); + if (clientIndex < 0) + { + _ = _requestStates.TryRemove(context, out _); + throw new InvalidOperationException("The invocation did not use a configured ordered failover client."); + } + + // Selection immediately removes this state before invoking the next client. + _requestStates[context] = new(clientIndex + 1, attempt.Exception); + + return default; + } + + /// + protected override void Dispose(bool disposing) + { + if (_disposed) + { + return; + } + + _disposed = true; + _requestStates.Clear(); + + if (disposing && !_leaveOpen) + { + foreach (IChatClient client in _clients) + { + client.Dispose(); + } + } + + base.Dispose(disposing); + } + + private int IndexOfClient(IChatClient client) + { + for (int i = 0; i < _clients.Length; i++) + { + if (ReferenceEquals(_clients[i], client)) + { + return i; + } + } + + return -1; + } + + private sealed class RequestState(int clientIndex, Exception lastException) + { + public int ClientIndex { get; } = clientIndex; + + public Exception LastException { get; } = lastException; + } +} diff --git a/src/Libraries/Microsoft.Extensions.AI/ChatRouting/SemanticRoutingChatClient.cs b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/SemanticRoutingChatClient.cs new file mode 100644 index 00000000000..f51aa5c33d7 --- /dev/null +++ b/src/Libraries/Microsoft.Extensions.AI/ChatRouting/SemanticRoutingChatClient.cs @@ -0,0 +1,360 @@ +// 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.Linq; +using System.Numerics.Tensors; +using System.Threading; +using System.Threading.Tasks; +using Microsoft.Shared.DiagnosticIds; +using Microsoft.Shared.Diagnostics; + +namespace Microsoft.Extensions.AI; + +/// Routes requests by semantic similarity to app-provided example utterances. +/// +/// +/// Profile embeddings are generated lazily and cached. Each request embeds the last user message and selects the +/// client with the highest score after aggregating the cosine similarities of the best-matching profile utterances. +/// The configured default client is selected when no user message is available or when the highest score is below the +/// configured threshold. +/// +/// +/// The configured client instances are used as stable routing identities and are distinguished by reference. +/// Per-call options do not participate in that identity. By default this instance owns the clients and embedding +/// generator and disposes them when it is disposed. +/// +/// +/// The example-utterance routing approach is inspired by +/// Aurelio Labs' semantic-router project. +/// +/// +[Experimental(DiagnosticIds.Experiments.AIRoutingChat, UrlFormat = DiagnosticIds.UrlFormat)] +public sealed class SemanticRoutingChatClient : RoutingChatClient +{ + /// Specifies how profile similarity scores are aggregated for each client. + public enum ScoreAggregation + { + /// Average the matching profile scores for each client. + Mean, + + /// Sum the matching profile scores for each client. + Sum, + } + + private readonly IChatClient[] _clients; + private readonly IEmbeddingGenerator> _embeddingGenerator; + private readonly SemaphoreSlim _indexGate = new(1, 1); + private readonly bool _leaveOpen; + private readonly (int ClientIndex, string Text)[] _profiles; + private readonly ScoreAggregation _scoreAggregation; + private readonly float _scoreThreshold; + private readonly int _topK; + + private bool _disposed; + private EmbeddedProfile[]? _index; + + /// Initializes a new instance of the class. + /// The generator used to embed profile utterances and request text. + /// The example utterances associated with each client. + /// + /// The client selected when no profile satisfies . + /// + /// The minimum aggregated score required to select a profiled client. + /// + /// The number of highest-scoring profile utterances, across all clients, whose scores are aggregated. + /// The default is 1. + /// + /// The method used to aggregate matching profile scores for each client. + /// + /// to leave the configured clients and embedding generator open when this instance is + /// disposed; otherwise, . The default is . + /// + /// + /// , , or + /// is . + /// + /// + /// is empty or contains a null client, an empty utterance list, or a blank + /// utterance. + /// + /// + /// is not positive, is invalid, or + /// is outside the possible range for the configured aggregation. + /// + public SemanticRoutingChatClient( + IEmbeddingGenerator> embeddingGenerator, + IReadOnlyDictionary> clientProfiles, + IChatClient defaultClient, + float scoreThreshold = 0.3f, + int topK = 1, + ScoreAggregation scoreAggregation = ScoreAggregation.Mean, + bool leaveOpen = false) + { + _ = Throw.IfNull(embeddingGenerator); + _ = Throw.IfNull(clientProfiles); + _ = Throw.IfNull(defaultClient); + + if (clientProfiles.Count == 0) + { + Throw.ArgumentException(nameof(clientProfiles), "At least one client profile must be provided."); + } + + if (topK <= 0) + { + Throw.ArgumentOutOfRangeException(nameof(topK)); + } + + if (scoreAggregation is not ScoreAggregation.Mean and not ScoreAggregation.Sum) + { + Throw.ArgumentOutOfRangeException(nameof(scoreAggregation)); + } + + float scoreLimit = scoreAggregation == ScoreAggregation.Sum ? topK : 1; + if (float.IsNaN(scoreThreshold) || + float.IsInfinity(scoreThreshold) || + scoreThreshold < -scoreLimit || + scoreThreshold > scoreLimit) + { + Throw.ArgumentOutOfRangeException(nameof(scoreThreshold)); + } + + _embeddingGenerator = embeddingGenerator; + _leaveOpen = leaveOpen; + _topK = topK; + _scoreAggregation = scoreAggregation; + _scoreThreshold = scoreThreshold; + + var profiles = new List<(int ClientIndex, string Text)>(); + var clients = new List { defaultClient }; + foreach (KeyValuePair> profile in clientProfiles) + { + IChatClient client = profile.Key; + IReadOnlyList utterances = profile.Value; + if (client is null) + { + Throw.ArgumentException(nameof(clientProfiles), "Profile clients must not be null."); + } + + if (utterances is null || utterances.Count == 0) + { + Throw.ArgumentException( + nameof(clientProfiles), + "Every profile client must have at least one example utterance."); + } + + int clientIndex = clients.FindIndex(candidate => ReferenceEquals(candidate, client)); + if (clientIndex < 0) + { + clientIndex = clients.Count; + clients.Add(client); + } + + foreach (string utterance in utterances) + { + if (string.IsNullOrWhiteSpace(utterance)) + { + Throw.ArgumentException(nameof(clientProfiles), "Profile utterances must not be blank."); + } + + profiles.Add((clientIndex, utterance)); + } + } + + _clients = [.. clients]; + _profiles = [.. profiles]; + } + + /// + protected override async ValueTask SelectClientAsync( + RoutingContext context, + CancellationToken cancellationToken) + { + _ = Throw.IfNull(context); + + string? query = context.Messages.LastOrDefault( + static message => message.Role == ChatRole.User)?.Text; + if (string.IsNullOrWhiteSpace(query)) + { + return new(_clients[0], context.ChatOptions); + } + + EmbeddedProfile[] index = await EnsureIndexAsync(cancellationToken).ConfigureAwait(false); + GeneratedEmbeddings> generated = + await _embeddingGenerator.GenerateAsync( + [query!], + cancellationToken: cancellationToken).ConfigureAwait(false) ?? + throw new InvalidOperationException("The embedding generator returned null."); + if (generated.Count != 1) + { + throw new InvalidOperationException("The embedding generator did not return one query embedding."); + } + + ReadOnlySpan queryVector = generated[0].Vector.Span; + if (queryVector.Length != index[0].Vector.Length) + { + throw new InvalidOperationException( + "The query embedding dimension does not match the profile embedding dimension."); + } + + if (_topK == 1) + { + int bestClientIndex = -1; + float bestScore = float.NegativeInfinity; + foreach (EmbeddedProfile profile in index) + { + float score = TensorPrimitives.CosineSimilarity(queryVector, profile.Vector); + if (score > bestScore) + { + bestClientIndex = profile.ClientIndex; + bestScore = score; + } + } + + IChatClient client = bestClientIndex >= 0 && bestScore >= _scoreThreshold + ? _clients[bestClientIndex] + : _clients[0]; + + return new(client, context.ChatOptions); + } + + var matches = new ScoredProfile[index.Length]; + for (int i = 0; i < index.Length; i++) + { + matches[i] = new( + i, + TensorPrimitives.CosineSimilarity(queryVector, index[i].Vector)); + } + + Array.Sort(matches, static (left, right) => + { + int scoreComparison = right.Score.CompareTo(left.Score); + return scoreComparison != 0 + ? scoreComparison + : left.ProfileIndex.CompareTo(right.ProfileIndex); + }); + + int matchCount = Math.Min(_topK, matches.Length); + var scoreSums = new float[_clients.Length]; + var scoreCounts = new int[_clients.Length]; + var clientOrder = new int[_clients.Length]; + int clientCount = 0; + for (int i = 0; i < matchCount; i++) + { + EmbeddedProfile profile = index[matches[i].ProfileIndex]; + int clientIndex = profile.ClientIndex; + if (scoreCounts[clientIndex] == 0) + { + clientOrder[clientCount++] = clientIndex; + } + + scoreSums[clientIndex] += matches[i].Score; + scoreCounts[clientIndex]++; + } + + int bestAggregatedClientIndex = -1; + float bestAggregatedScore = float.NegativeInfinity; + for (int i = 0; i < clientCount; i++) + { + int clientIndex = clientOrder[i]; + float score = _scoreAggregation == ScoreAggregation.Mean + ? scoreSums[clientIndex] / scoreCounts[clientIndex] + : scoreSums[clientIndex]; + if (score > bestAggregatedScore) + { + bestAggregatedClientIndex = clientIndex; + bestAggregatedScore = score; + } + } + + IChatClient selectedClient = bestAggregatedClientIndex >= 0 && bestAggregatedScore >= _scoreThreshold + ? _clients[bestAggregatedClientIndex] + : _clients[0]; + + return new(selectedClient, context.ChatOptions); + } + + /// + protected override void Dispose(bool disposing) + { + if (_disposed) + { + return; + } + + _disposed = true; + if (disposing) + { + _indexGate.Dispose(); + if (!_leaveOpen) + { + foreach (IChatClient client in _clients) + { + client.Dispose(); + } + + _embeddingGenerator.Dispose(); + } + } + + base.Dispose(disposing); + } + + private async Task EnsureIndexAsync(CancellationToken cancellationToken) + { + if (_index is { } cached) + { + return cached; + } + + await _indexGate.WaitAsync(cancellationToken).ConfigureAwait(false); + try + { + if (_index is { } existing) + { + return existing; + } + + GeneratedEmbeddings> embeddings = + await _embeddingGenerator.GenerateAsync( + _profiles.Select(profile => profile.Text), + cancellationToken: cancellationToken).ConfigureAwait(false) ?? + throw new InvalidOperationException("The embedding generator returned null."); + if (embeddings.Count != _profiles.Length) + { + throw new InvalidOperationException( + "The embedding generator did not return one embedding per profile utterance."); + } + + int dimensions = embeddings[0].Vector.Length; + if (dimensions == 0) + { + throw new InvalidOperationException("Profile embeddings must not be empty."); + } + + var index = new EmbeddedProfile[_profiles.Length]; + for (int i = 0; i < index.Length; i++) + { + if (embeddings[i].Vector.Length != dimensions) + { + throw new InvalidOperationException( + "All profile embeddings must have the same dimension."); + } + + index[i] = new(_profiles[i].ClientIndex, embeddings[i].Vector.ToArray()); + } + + return _index = index; + } + finally + { + _ = _indexGate.Release(); + } + } + + private sealed record EmbeddedProfile(int ClientIndex, float[] Vector); + + private readonly record struct ScoredProfile(int ProfileIndex, float Score); +} diff --git a/src/Libraries/Microsoft.Extensions.AI/Microsoft.Extensions.AI.json b/src/Libraries/Microsoft.Extensions.AI/Microsoft.Extensions.AI.json index 9efb90882b9..1dbff0278a5 100644 --- a/src/Libraries/Microsoft.Extensions.AI/Microsoft.Extensions.AI.json +++ b/src/Libraries/Microsoft.Extensions.AI/Microsoft.Extensions.AI.json @@ -489,6 +489,68 @@ } ] }, + { + "Type": "abstract class Microsoft.Extensions.AI.FailoverChatClient : Microsoft.Extensions.AI.RoutingChatClient", + "Stage": "Experimental", + "Methods": [ + { + "Member": "Microsoft.Extensions.AI.FailoverChatClient.FailoverChatClient();", + "Stage": "Experimental" + }, + { + "Member": "sealed override System.Threading.Tasks.Task Microsoft.Extensions.AI.FailoverChatClient.GetResponseAsync(System.Collections.Generic.IEnumerable messages, Microsoft.Extensions.AI.ChatOptions? options = null, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken));", + "Stage": "Experimental" + }, + { + "Member": "sealed override System.Collections.Generic.IAsyncEnumerable Microsoft.Extensions.AI.FailoverChatClient.GetStreamingResponseAsync(System.Collections.Generic.IEnumerable messages, Microsoft.Extensions.AI.ChatOptions? options = null, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken));", + "Stage": "Experimental" + }, + { + "Member": "virtual System.Threading.Tasks.ValueTask Microsoft.Extensions.AI.FailoverChatClient.OnRoutingUpdateAsync(Microsoft.Extensions.AI.RoutingContext context, Microsoft.Extensions.AI.FailoverChatClientAttempt attempt, bool isTerminal, System.Threading.CancellationToken cancellationToken);", + "Stage": "Experimental" + } + ], + "Properties": [ + { + "Member": "int? Microsoft.Extensions.AI.FailoverChatClient.MaximumAttemptsPerRequest { get; set; }", + "Stage": "Experimental" + } + ] + }, + { + "Type": "sealed class Microsoft.Extensions.AI.FailoverChatClientAttempt", + "Stage": "Experimental", + "Properties": [ + { + "Member": "Microsoft.Extensions.AI.ChatOptions? Microsoft.Extensions.AI.FailoverChatClientAttempt.ChatOptions { get; }", + "Stage": "Experimental" + }, + { + "Member": "Microsoft.Extensions.AI.IChatClient Microsoft.Extensions.AI.FailoverChatClientAttempt.Client { get; }", + "Stage": "Experimental" + }, + { + "Member": "System.TimeSpan Microsoft.Extensions.AI.FailoverChatClientAttempt.Duration { get; }", + "Stage": "Experimental" + }, + { + "Member": "System.Exception? Microsoft.Extensions.AI.FailoverChatClientAttempt.Exception { get; }", + "Stage": "Experimental" + }, + { + "Member": "bool Microsoft.Extensions.AI.FailoverChatClientAttempt.OutputCommitted { get; }", + "Stage": "Experimental" + }, + { + "Member": "bool Microsoft.Extensions.AI.FailoverChatClientAttempt.ResponseCompleted { get; }", + "Stage": "Experimental" + }, + { + "Member": "System.TimeSpan? Microsoft.Extensions.AI.FailoverChatClientAttempt.TimeToFirstUpdate { get; }", + "Stage": "Experimental" + } + ] + }, { "Type": "class Microsoft.Extensions.AI.FunctionInvocationContext", "Stage": "Stable", @@ -1425,6 +1487,28 @@ } ] }, + { + "Type": "sealed class Microsoft.Extensions.AI.OrderedFailoverChatClient : Microsoft.Extensions.AI.FailoverChatClient", + "Stage": "Experimental", + "Methods": [ + { + "Member": "Microsoft.Extensions.AI.OrderedFailoverChatClient.OrderedFailoverChatClient(System.Collections.Generic.IReadOnlyList clients, bool leaveOpen = false);", + "Stage": "Experimental" + }, + { + "Member": "override void Microsoft.Extensions.AI.OrderedFailoverChatClient.Dispose(bool disposing);", + "Stage": "Experimental" + }, + { + "Member": "override System.Threading.Tasks.ValueTask Microsoft.Extensions.AI.OrderedFailoverChatClient.OnRoutingUpdateAsync(Microsoft.Extensions.AI.RoutingContext context, Microsoft.Extensions.AI.FailoverChatClientAttempt attempt, bool isTerminal, System.Threading.CancellationToken cancellationToken);", + "Stage": "Experimental" + }, + { + "Member": "override System.Threading.Tasks.ValueTask Microsoft.Extensions.AI.OrderedFailoverChatClient.SelectClientAsync(Microsoft.Extensions.AI.RoutingContext context, System.Threading.CancellationToken cancellationToken);", + "Stage": "Experimental" + } + ] + }, { "Type": "sealed class Microsoft.Extensions.AI.RealtimeClientBuilder", "Stage": "Experimental", @@ -1525,6 +1609,46 @@ } ] }, + { + "Type": "sealed class Microsoft.Extensions.AI.SemanticRoutingChatClient : Microsoft.Extensions.AI.RoutingChatClient", + "Stage": "Experimental", + "Methods": [ + { + "Member": "Microsoft.Extensions.AI.SemanticRoutingChatClient.SemanticRoutingChatClient(Microsoft.Extensions.AI.IEmbeddingGenerator> embeddingGenerator, System.Collections.Generic.IReadOnlyDictionary> clientProfiles, Microsoft.Extensions.AI.IChatClient defaultClient, float scoreThreshold = 0.3f, int topK = 1, Microsoft.Extensions.AI.SemanticRoutingChatClient.ScoreAggregation scoreAggregation = Microsoft.Extensions.AI.SemanticRoutingChatClient.ScoreAggregation.Mean, bool leaveOpen = false);", + "Stage": "Experimental" + }, + { + "Member": "override void Microsoft.Extensions.AI.SemanticRoutingChatClient.Dispose(bool disposing);", + "Stage": "Experimental" + }, + { + "Member": "override System.Threading.Tasks.ValueTask Microsoft.Extensions.AI.SemanticRoutingChatClient.SelectClientAsync(Microsoft.Extensions.AI.RoutingContext context, System.Threading.CancellationToken cancellationToken);", + "Stage": "Experimental" + } + ] + }, + { + "Type": "enum Microsoft.Extensions.AI.SemanticRoutingChatClient.ScoreAggregation", + "Stage": "Experimental", + "Methods": [ + { + "Member": "Microsoft.Extensions.AI.SemanticRoutingChatClient.ScoreAggregation.ScoreAggregation();", + "Stage": "Experimental" + } + ], + "Fields": [ + { + "Member": "const Microsoft.Extensions.AI.SemanticRoutingChatClient.ScoreAggregation Microsoft.Extensions.AI.SemanticRoutingChatClient.ScoreAggregation.Mean", + "Stage": "Experimental", + "Value": "0" + }, + { + "Member": "const Microsoft.Extensions.AI.SemanticRoutingChatClient.ScoreAggregation Microsoft.Extensions.AI.SemanticRoutingChatClient.ScoreAggregation.Sum", + "Stage": "Experimental", + "Value": "1" + } + ] + }, { "Type": "sealed class Microsoft.Extensions.AI.SpeechToTextClientBuilder", "Stage": "Experimental", diff --git a/src/Shared/DiagnosticIds/DiagnosticIds.cs b/src/Shared/DiagnosticIds/DiagnosticIds.cs index abddbf61c5a..29a3d276b48 100644 --- a/src/Shared/DiagnosticIds/DiagnosticIds.cs +++ b/src/Shared/DiagnosticIds/DiagnosticIds.cs @@ -61,6 +61,7 @@ internal static class Experiments internal const string AIRealTime = AIExperiments; internal const string AIFiles = AIExperiments; internal const string AIOpenAIRequestPolicies = AIExperiments; + internal const string AIRoutingChat = AIExperiments; // These diagnostic IDs are defined by the OpenAI package for its experimental APIs. // We use the same IDs so consumers do not need to suppress additional diagnostics diff --git a/test/Libraries/Microsoft.Extensions.AI.Abstractions.Tests/ChatRouting/RoutingChatClientTests.cs b/test/Libraries/Microsoft.Extensions.AI.Abstractions.Tests/ChatRouting/RoutingChatClientTests.cs new file mode 100644 index 00000000000..f15f3896df8 --- /dev/null +++ b/test/Libraries/Microsoft.Extensions.AI.Abstractions.Tests/ChatRouting/RoutingChatClientTests.cs @@ -0,0 +1,199 @@ +// 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.Threading; +using System.Threading.Tasks; +using Xunit; + +namespace Microsoft.Extensions.AI; + +public class RoutingChatClientTests +{ + [Fact] + public void Create_RejectsNullSelector() + { + Assert.Throws(() => RoutingChatClient.Create(null!)); + } + + [Fact] + public void RoutingSelection_ValidatesClientAndClonesOptions() + { + using var client = new TestChatClient(); + var options = new ChatOptions { ModelId = "original" }; + + var selection = new RoutingSelection(client, options); + options.ModelId = "changed"; + + Assert.Same(client, selection.Client); + Assert.NotSame(options, selection.ChatOptions); + Assert.Equal("original", selection.ChatOptions!.ModelId); + Assert.Throws(() => new RoutingSelection(null!, options)); + } + + [Fact] + public async Task Create_SelectsClientForRequest() + { + var messages = new ChatMessage[] { new(ChatRole.User, "hi") }; + var options = new ChatOptions { ModelId = "request" }; + using var cancellationSource = new CancellationTokenSource(); + ChatResponse expected = new(new ChatMessage(ChatRole.Assistant, "ok")); + ChatOptions? forwardedOptions = null; + using var selected = new TestChatClient + { + GetResponseAsyncCallback = (_, selectedOptions, _) => + { + forwardedOptions = selectedOptions; + return Task.FromResult(expected); + }, + }; + RoutingContext? observedContext = null; + CancellationToken observedToken = default; + int selectionCount = 0; + using RoutingChatClient router = RoutingChatClient.Create((context, cancellationToken) => + { + observedContext = context; + observedToken = cancellationToken; + selectionCount++; + ChatOptions selectedOptions = context.ChatOptions!.Clone(); + selectedOptions.ModelId = "selected"; + return new(new RoutingSelection(selected, selectedOptions)); + }); + + ChatResponse response = await router.GetResponseAsync(messages, options, cancellationSource.Token); + + Assert.Same(expected, response); + Assert.Same(messages, observedContext!.Messages); + Assert.NotSame(options, observedContext.ChatOptions); + Assert.NotSame(observedContext.ChatOptions, forwardedOptions); + Assert.Equal("request", observedContext.ChatOptions!.ModelId); + Assert.Equal("selected", forwardedOptions!.ModelId); + Assert.Equal("request", options.ModelId); + Assert.Equal(cancellationSource.Token, observedToken); + Assert.Equal(1, selectionCount); + } + + [Fact] + public async Task Create_DoesNotReselectAfterFailure() + { + var expected = new InvalidOperationException("failed"); + using var selected = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => throw expected, + }; + int selectionCount = 0; + using RoutingChatClient router = RoutingChatClient.Create((context, _) => + { + selectionCount++; + return new(new RoutingSelection(selected, context.ChatOptions)); + }); + + InvalidOperationException actual = await Assert.ThrowsAsync( + () => router.GetResponseAsync([new(ChatRole.User, "hi")])); + + Assert.Same(expected, actual); + Assert.Equal(1, selectionCount); + } + + [Fact] + public async Task Create_NullSelectionThrowsForNonStreamingAndStreaming() + { + using RoutingChatClient router = RoutingChatClient.Create((_, _) => new((RoutingSelection)null!)); + + InvalidOperationException nonStreaming = await Assert.ThrowsAsync( + () => router.GetResponseAsync([new(ChatRole.User, "hi")])); + InvalidOperationException streaming = await Assert.ThrowsAsync( + () => CollectAsync(router.GetStreamingResponseAsync([new(ChatRole.User, "hi")]))); + + Assert.Contains("SelectClientAsync", nonStreaming.Message, StringComparison.Ordinal); + Assert.Contains("SelectClientAsync", streaming.Message, StringComparison.Ordinal); + } + + [Fact] + public async Task Create_RejectsNullMessagesForNonStreamingAndStreaming() + { + using var selected = new TestChatClient(); + using RoutingChatClient router = RoutingChatClient.Create( + (context, _) => new(new RoutingSelection(selected, context.ChatOptions))); + + await Assert.ThrowsAsync(() => router.GetResponseAsync(null!)); + await Assert.ThrowsAsync( + () => CollectAsync(router.GetStreamingResponseAsync(null!))); + } + + [Fact] + public async Task Create_DoesNotDisposeSelectedClient() + { + using var selected = new CountingDisposeClient(); + using (RoutingChatClient router = RoutingChatClient.Create( + (context, _) => new(new RoutingSelection(selected, context.ChatOptions)))) + { + _ = await router.GetResponseAsync([new(ChatRole.User, "hi")]); + } + + Assert.Equal(0, selected.DisposeCount); + } + + [Fact] + public void GetService_ReturnsSelfAndNullForUnknownOrKeyed() + { + using var client = new DelegatingTestRouter(_ => throw new NotSupportedException()); + + Assert.Same(client, client.GetService(typeof(DelegatingTestRouter))); + Assert.Same(client, client.GetService(typeof(RoutingChatClient))); + Assert.Same(client, client.GetService(typeof(IChatClient))); + Assert.Null(client.GetService(typeof(DelegatingTestRouter), serviceKey: "key")); + Assert.Null(client.GetService(typeof(string))); + } + + private static async Task> CollectAsync( + IAsyncEnumerable updates) + { + var result = new List(); + await foreach (ChatResponseUpdate update in updates) + { + result.Add(update); + } + + return result; + } + + private sealed class DelegatingTestRouter : RoutingChatClient + { + private readonly Func _select; + + public DelegatingTestRouter(Func select) + { + _select = select; + } + + protected override ValueTask SelectClientAsync( + RoutingContext context, + CancellationToken cancellationToken) => + new(new RoutingSelection(_select(context), context.ChatOptions)); + } + + private sealed class CountingDisposeClient : IChatClient + { + public int DisposeCount { get; private set; } + + public Task GetResponseAsync( + IEnumerable messages, ChatOptions? options = null, + CancellationToken cancellationToken = default) => + Task.FromResult(new ChatResponse()); + + public IAsyncEnumerable GetStreamingResponseAsync( + IEnumerable messages, ChatOptions? options = null, + CancellationToken cancellationToken = default) => + throw new NotSupportedException(); + + public object? GetService(Type serviceType, object? serviceKey = null) => null; + + public void Dispose() => DisposeCount++; + + public override bool Equals(object? obj) => obj is CountingDisposeClient; + + public override int GetHashCode() => 0; + } +} diff --git a/test/Libraries/Microsoft.Extensions.AI.Abstractions.Tests/ChatRouting/RoutingContextTests.cs b/test/Libraries/Microsoft.Extensions.AI.Abstractions.Tests/ChatRouting/RoutingContextTests.cs new file mode 100644 index 00000000000..d97195262a7 --- /dev/null +++ b/test/Libraries/Microsoft.Extensions.AI.Abstractions.Tests/ChatRouting/RoutingContextTests.cs @@ -0,0 +1,26 @@ +// 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 Xunit; + +namespace Microsoft.Extensions.AI; + +public class RoutingContextTests +{ + [Fact] + public void RoutingContext_CarriesRequestInputs() + { + var messages = new List { new(ChatRole.User, "initial") }; + var options = new ChatOptions { ModelId = "initial" }; + var context = new RoutingContext(messages, options); + + Assert.Same(messages, context.Messages); + Assert.NotSame(options, context.ChatOptions); + Assert.Equal("initial", context.ChatOptions!.ModelId); + context.ChatOptions.ModelId = "changed"; + Assert.Equal("initial", options.ModelId); + Assert.Throws(() => new RoutingContext(null!, options)); + } +} diff --git a/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/FailoverChatClientTests.cs b/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/FailoverChatClientTests.cs new file mode 100644 index 00000000000..b2c0289c345 --- /dev/null +++ b/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/FailoverChatClientTests.cs @@ -0,0 +1,1450 @@ +// 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.Runtime.CompilerServices; +using System.Threading; +using System.Threading.Tasks; +using Xunit; + +namespace Microsoft.Extensions.AI; + +public class FailoverChatClientTests +{ + [Fact] + public void MaximumAttemptsPerRequest_ValidatesValue() + { + using var client = new DelegatingFailoverTestRouter( + _ => throw new NotSupportedException()); + + Assert.Null(client.MaximumAttemptsPerRequest); + client.MaximumAttemptsPerRequest = 2; + Assert.Equal(2, client.MaximumAttemptsPerRequest); + client.MaximumAttemptsPerRequest = null; + Assert.Null(client.MaximumAttemptsPerRequest); + Assert.Throws(() => client.MaximumAttemptsPerRequest = 0); + Assert.Throws(() => client.MaximumAttemptsPerRequest = -1); + } + + [Fact] + public async Task Failover_RejectsNullMessagesForNonStreamingAndStreaming() + { + using var selected = new TestChatClient(); + using var router = new DelegatingFailoverTestRouter(_ => selected); + + await Assert.ThrowsAsync(() => router.GetResponseAsync(null!)); + await Assert.ThrowsAsync( + () => CollectAsync(router.GetStreamingResponseAsync(null!))); + } + + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task InitialSelectionFailureDoesNotReportAttempt(bool streaming) + { + var expected = new InvalidOperationException("selection failed"); + RoutingContext? selectedContext = null; + int updateCount = 0; + using var router = new DelegatingFailoverTestRouter( + context => + { + selectedContext = context; + throw expected; + }, + (_, _, _) => updateCount++); + + Task operation = streaming + ? router.GetStreamingResponseAsync([new(ChatRole.User, "hi")]).ToChatResponseAsync() + : router.GetResponseAsync([new(ChatRole.User, "hi")]); + InvalidOperationException actual = + await Assert.ThrowsAsync(() => operation); + + Assert.Same(expected, actual); + Assert.NotNull(selectedContext); + Assert.Equal(0, updateCount); + } + + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task InitialSelectionFailureWithCallerCancellationPropagatesSelectionFailure(bool streaming) + { + using var cancellationSource = new CancellationTokenSource(); + cancellationSource.Cancel(); + var selectionException = new InvalidOperationException("selection failed"); + int updateCount = 0; + using var router = new DelegatingFailoverTestRouter( + _ => throw selectionException, + (_, _, _) => updateCount++); + + Task operation = streaming + ? router.GetStreamingResponseAsync( + [new(ChatRole.User, "hi")], + cancellationToken: cancellationSource.Token).ToChatResponseAsync() + : router.GetResponseAsync( + [new(ChatRole.User, "hi")], + cancellationToken: cancellationSource.Token); + InvalidOperationException actual = + await Assert.ThrowsAsync(() => operation); + + Assert.Same(selectionException, actual); + Assert.Equal(0, updateCount); + } + + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task NullSelectionDoesNotReportAttempt(bool streaming) + { + int updateCount = 0; + using var router = new DelegatingFailoverTestRouter( + _ => null!, + (_, _, _) => updateCount++); + + Task operation = streaming + ? router.GetStreamingResponseAsync([new(ChatRole.User, "hi")]).ToChatResponseAsync() + : router.GetResponseAsync([new(ChatRole.User, "hi")]); + InvalidOperationException exception = + await Assert.ThrowsAsync(() => operation); + + Assert.Contains("SelectClientAsync", exception.Message, StringComparison.Ordinal); + Assert.Equal(0, updateCount); + } + + [Fact] + public async Task Failover_RetrySelectionFailureReportsOnlyCompletedAttempt() + { + var invocationException = new InvalidOperationException("invocation failed"); + var selectionException = new InvalidOperationException("selection failed"); + using var failing = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => throw invocationException, + }; + int selections = 0; + var updates = new List<(FailoverChatClientAttempt attempt, bool isTerminal)>(); + using var router = new DelegatingFailoverTestRouter( + _ => ++selections == 1 ? failing : throw selectionException, + (_, attempt, isTerminal) => + { + updates.Add((attempt, isTerminal)); + }); + + InvalidOperationException actual = await Assert.ThrowsAsync( + () => router.GetResponseAsync([new(ChatRole.User, "hi")])); + + Assert.Same(selectionException, actual); + Assert.Equal(2, selections); + (FailoverChatClientAttempt attempt, bool isTerminal) update = Assert.Single(updates); + Assert.Same(failing, update.attempt.Client); + Assert.Same(invocationException, update.attempt.Exception); + Assert.False(update.isTerminal); + } + + [Fact] + public async Task TerminalRoutingUpdateFailureReplacesRequestFailure() + { + var invocationException = new InvalidOperationException("invocation failed"); + var completionException = new InvalidOperationException("completion failed"); + using var failing = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => throw invocationException, + }; + using var router = new DelegatingFailoverTestRouter( + _ => failing, + (_, _, isTerminal) => + { + Assert.True(isTerminal); + throw completionException; + }) + { + MaximumAttemptsPerRequest = 1, + }; + + InvalidOperationException actual = await Assert.ThrowsAsync( + () => router.GetResponseAsync([new(ChatRole.User, "hi")])); + + Assert.Same(completionException, actual); + } + + [Fact] + public async Task NonterminalRoutingUpdateFailureStopsWithoutAnotherUpdate() + { + var invocationException = new InvalidOperationException("invocation failed"); + var updateException = new InvalidOperationException("update failed"); + using var failing = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => throw invocationException, + }; + int selections = 0; + int updates = 0; + using var router = new DelegatingFailoverTestRouter( + _ => + { + selections++; + return failing; + }, + (_, attempt, isTerminal) => + { + updates++; + Assert.Same(invocationException, attempt!.Exception); + Assert.False(isTerminal); + throw updateException; + }); + + InvalidOperationException actual = await Assert.ThrowsAsync( + () => router.GetResponseAsync([new(ChatRole.User, "hi")])); + + Assert.Same(updateException, actual); + Assert.Equal(1, selections); + Assert.Equal(1, updates); + } + + [Fact] + public async Task Dispatch_ConfiguredClientPreservesAndOverridesRequestOptions() + { + ChatOptions? forwarded = null; + using var inner = new TestChatClient + { + GetResponseAsyncCallback = (_, options, _) => + { + forwarded = options; + return Task.FromResult(new ChatResponse()); + }, + }; + using var configured = new ConfigureOptionsChatClient( + inner, + options => options.ModelId = "route"); + var requestOptions = new ChatOptions + { + Instructions = "caller", + ModelId = "request", + }; + FailoverChatClientAttempt? observedAttempt = null; + using var router = new DelegatingFailoverTestRouter( + _ => configured, + (_, attempt, isTerminal) => + { + Assert.True(isTerminal); + observedAttempt = attempt; + }); + + _ = await router.GetResponseAsync([new(ChatRole.User, "hi")], requestOptions); + + Assert.NotSame(requestOptions, forwarded); + Assert.Equal("route", forwarded!.ModelId); + Assert.Equal("caller", forwarded.Instructions); + Assert.Equal("request", observedAttempt!.ChatOptions!.ModelId); + Assert.Equal("request", requestOptions.ModelId); + } + + [Fact] + public async Task Failover_UsesFreshRequestOptionsForEachAttempt() + { + var requestOptions = new ChatOptions + { + Instructions = "caller", + ModelId = "request", + }; + var invokedOptions = new List(); + using var first = new TestChatClient + { + GetResponseAsyncCallback = (_, options, _) => + { + invokedOptions.Add(options!); + options!.ModelId = "changed by first client"; + throw new InvalidOperationException("failed"); + }, + }; + ChatResponse expected = new(new ChatMessage(ChatRole.Assistant, "ok")); + using var second = new TestChatClient + { + GetResponseAsyncCallback = (_, options, _) => + { + invokedOptions.Add(options!); + return Task.FromResult(expected); + }, + }; + int selections = 0; + using var router = new DelegatingFailoverTestRouter( + _ => ++selections == 1 ? first : second); + + ChatResponse response = await router.GetResponseAsync( + [new(ChatRole.User, "hi")], + requestOptions); + + Assert.Same(expected, response); + Assert.Equal(2, selections); + Assert.Equal(2, invokedOptions.Count); + Assert.NotSame(invokedOptions[0], invokedOptions[1]); + Assert.Equal("changed by first client", invokedOptions[0].ModelId); + Assert.Equal("request", invokedOptions[1].ModelId); + Assert.Equal("caller", invokedOptions[1].Instructions); + Assert.Equal("request", requestOptions.ModelId); + } + + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task SelectedOptionsAreScopedToAttempt(bool streaming) + { + var requestOptions = new ChatOptions + { + Instructions = "caller", + ModelId = "request", + }; + var invokedOptions = new List(); + var attempts = new List(); + using var first = new TestChatClient + { + GetResponseAsyncCallback = (_, options, _) => + { + invokedOptions.Add(options!); + options!.ModelId = "mutated"; + throw new InvalidOperationException("failed"); + }, + GetStreamingResponseAsyncCallback = (_, options, _) => + { + invokedOptions.Add(options!); + options!.ModelId = "mutated"; + return ThrowingStream("failed"); + }, + }; + ChatResponse expected = new(new ChatMessage(ChatRole.Assistant, "ok")); + using var second = new TestChatClient + { + GetResponseAsyncCallback = (_, options, _) => + { + invokedOptions.Add(options!); + return Task.FromResult(expected); + }, + GetStreamingResponseAsyncCallback = (_, options, _) => + { + invokedOptions.Add(options!); + return YieldUpdates("ok"); + }, + }; + int selections = 0; + using var router = new DelegatingSelectionFailoverTestRouter( + context => + { + Assert.Equal("request", context.ChatOptions!.ModelId); + ChatOptions selectedOptions = context.ChatOptions.Clone(); + bool firstAttempt = selections++ == 0; + selectedOptions.ModelId = firstAttempt ? "first" : "second"; + return new(firstAttempt ? first : second, selectedOptions); + }, + (_, attempt, _) => attempts.Add(attempt)); + + ChatResponse response = streaming + ? await router.GetStreamingResponseAsync( + [new(ChatRole.User, "hi")], + requestOptions).ToChatResponseAsync() + : await router.GetResponseAsync([new(ChatRole.User, "hi")], requestOptions); + + Assert.Equal("ok", response.Text); + Assert.Equal(2, selections); + Assert.Equal("mutated", invokedOptions[0].ModelId); + Assert.Equal("second", invokedOptions[1].ModelId); + Assert.Collection( + attempts, + attempt => + { + Assert.Equal("first", attempt.ChatOptions!.ModelId); + Assert.NotSame(invokedOptions[0], attempt.ChatOptions); + }, + attempt => + { + Assert.Equal("second", attempt.ChatOptions!.ModelId); + Assert.NotSame(invokedOptions[1], attempt.ChatOptions); + }); + Assert.Equal("request", requestOptions.ModelId); + } + + [Fact] + public async Task Failure_Propagates() + { + var expected = new InvalidOperationException("failed"); + using var inner = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => throw expected, + }; + FailoverChatClientAttempt? observed = null; + using var router = new DelegatingFailoverTestRouter( + _ => inner, + (_, attempt, isTerminal) => + { + Assert.True(isTerminal); + observed = attempt; + }) + { + MaximumAttemptsPerRequest = 1, + }; + + InvalidOperationException actual = await Assert.ThrowsAsync( + () => router.GetResponseAsync([new(ChatRole.User, "hi")])); + + Assert.Same(expected, actual); + Assert.Same(inner, observed!.Client); + Assert.Same(expected, observed.Exception); + Assert.False(observed.ResponseCompleted); + } + + [Fact] + public async Task Failover_UpdateChangesStateBeforeNextSelection() + { + using var failing = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => throw new InvalidOperationException("failed"), + }; + ChatResponse expected = new(new ChatMessage(ChatRole.Assistant, "ok")); + using var working = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => Task.FromResult(expected), + }; + RoutingContext? initialContext = null; + FailoverChatClientAttempt? failedAttempt = null; + FailoverChatClientAttempt? terminalAttempt = null; + using var router = new DelegatingFailoverTestRouter( + context => + { + initialContext ??= context; + Assert.Same(initialContext, context); + return failedAttempt is null ? failing : working; + }, + (context, attempt, isTerminal) => + { + Assert.Same(initialContext, context); + if (isTerminal) + { + terminalAttempt = attempt; + } + else + { + failedAttempt = attempt; + } + }); + + ChatResponse response = await router.GetResponseAsync([new(ChatRole.User, "hi")]); + + Assert.Same(expected, response); + Assert.Same(failing, failedAttempt!.Client); + Assert.Equal("failed", Assert.IsType(failedAttempt.Exception).Message); + Assert.True(failedAttempt.Duration >= TimeSpan.Zero); + Assert.Null(failedAttempt.TimeToFirstUpdate); + Assert.False(failedAttempt.OutputCommitted); + Assert.False(failedAttempt.ResponseCompleted); + Assert.Same(working, terminalAttempt!.Client); + Assert.Null(terminalAttempt.Exception); + Assert.True(terminalAttempt.ResponseCompleted); + } + + [Fact] + public async Task Failover_CanRetrySameClient() + { + int calls = 0; + using var inner = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => + ++calls == 1 + ? throw new InvalidOperationException("transient") + : Task.FromResult(new ChatResponse()), + }; + using var router = new DelegatingFailoverTestRouter( + _ => inner); + + _ = await router.GetResponseAsync([new(ChatRole.User, "hi")]); + + Assert.Equal(2, calls); + } + + [Fact] + public async Task MaximumAttemptsPerRequest_AllowsSuccessAtLimit() + { + int calls = 0; + ChatResponse expected = new(new ChatMessage(ChatRole.Assistant, "ok")); + using var inner = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => + ++calls == 1 + ? throw new InvalidOperationException("transient") + : Task.FromResult(expected), + }; + using var router = new DelegatingFailoverTestRouter( + _ => inner) + { + MaximumAttemptsPerRequest = 2, + }; + + ChatResponse response = await router.GetResponseAsync([new(ChatRole.User, "hi")]); + + Assert.Same(expected, response); + Assert.Equal(2, calls); + } + + [Fact] + public async Task MaximumAttemptsPerRequest_RethrowsLastFailure() + { + int calls = 0; + var terminalUpdates = new List(); + using var inner = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => + throw new InvalidOperationException($"failure {++calls}"), + }; + using var router = new DelegatingFailoverTestRouter( + _ => inner, + (_, attempt, isTerminal) => + { + Assert.NotNull(attempt); + terminalUpdates.Add(isTerminal); + }) + { + MaximumAttemptsPerRequest = 2, + }; + + InvalidOperationException exception = await Assert.ThrowsAsync( + () => router.GetResponseAsync([new(ChatRole.User, "hi")])); + + Assert.Equal("failure 2", exception.Message); + Assert.Equal(2, calls); + Assert.Equal([false, true], terminalUpdates); + } + + [Fact] + public async Task Failover_NullResponseIsReturnedWithoutFailover() + { + using var nullClient = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => Task.FromResult(null!), + }; + int selections = 0; + FailoverChatClientAttempt? terminalAttempt = null; + using var router = new DelegatingFailoverTestRouter( + _ => + { + selections++; + return nullClient; + }, + (_, attempt, isTerminal) => + { + Assert.True(isTerminal); + terminalAttempt = attempt; + }); + + ChatResponse? response = await router.GetResponseAsync([new(ChatRole.User, "hi")]); + + Assert.Null(response); + Assert.Equal(1, selections); + Assert.Null(terminalAttempt!.Exception); + Assert.True(terminalAttempt!.ResponseCompleted); + } + + [Fact] + public async Task Cancellation_DoesNotReselect() + { + using var cancellationSource = new CancellationTokenSource(); + cancellationSource.Cancel(); + int selections = 0; + using var inner = new TestChatClient + { + GetResponseAsyncCallback = (_, _, cancellationToken) => + throw new OperationCanceledException(cancellationToken), + }; + using var router = new DelegatingTestRouter(_ => + { + selections++; + return inner; + }); + + await Assert.ThrowsAnyAsync( + () => router.GetResponseAsync( + [new(ChatRole.User, "hi")], + cancellationToken: cancellationSource.Token)); + + Assert.Equal(1, selections); + } + + [Fact] + public async Task Failover_CancellationReportsTerminalUpdate() + { + using var cancellationSource = new CancellationTokenSource(); + cancellationSource.Cancel(); + var failure = new InvalidOperationException("failed"); + using var inner = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => throw failure, + }; + int updateCount = 0; + using var router = new DelegatingFailoverTestRouter( + _ => inner, + (_, attempt, isTerminal) => + { + updateCount++; + Assert.Same(failure, attempt!.Exception); + Assert.True(isTerminal); + }); + + await Assert.ThrowsAnyAsync( + () => router.GetResponseAsync( + [new(ChatRole.User, "hi")], + cancellationToken: cancellationSource.Token)); + + Assert.Equal(1, updateCount); + } + + [Fact] + public async Task Failover_CancellationDuringUpdateIsObservedByNextAttempt() + { + using var cancellationSource = new CancellationTokenSource(); + using var inner = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => throw new InvalidOperationException("failed"), + }; + int selections = 0; + var updates = new List<(FailoverChatClientAttempt attempt, bool isTerminal)>(); + using var router = new DelegatingFailoverTestRouter( + _ => + { + selections++; + return inner; + }, + (_, attempt, isTerminal) => + { + updates.Add((attempt, isTerminal)); + if (!isTerminal) + { + cancellationSource.Cancel(); + } + }); + + await Assert.ThrowsAnyAsync( + () => router.GetResponseAsync( + [new(ChatRole.User, "hi")], + cancellationToken: cancellationSource.Token)); + + Assert.Equal(2, selections); + Assert.Collection( + updates, + update => + { + Assert.NotNull(update.attempt); + Assert.False(update.isTerminal); + }, + update => + { + Assert.NotNull(update.attempt); + Assert.True(update.isTerminal); + }); + } + + [Fact] + public async Task Streaming_PreOutputFailurePropagates() + { + using var inner = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, _, _) => ThrowingStream("failed"), + }; + FailoverChatClientAttempt? observed = null; + using var router = new DelegatingFailoverTestRouter( + _ => inner, + (_, attempt, isTerminal) => + { + Assert.True(isTerminal); + observed = attempt; + }) + { + MaximumAttemptsPerRequest = 1, + }; + + InvalidOperationException exception = await Assert.ThrowsAsync( + () => CollectAsync(router.GetStreamingResponseAsync([new(ChatRole.User, "hi")]))); + + Assert.Equal("failed", exception.Message); + Assert.Same(exception, observed!.Exception); + Assert.False(observed.OutputCommitted); + Assert.False(observed.ResponseCompleted); + } + + [Fact] + public async Task Streaming_FallsBackBeforeFirstUpdate() + { + using var failing = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, _, _) => ThrowingStream("failed"), + }; + using var working = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, _, _) => YieldUpdates("ok"), + }; + int selections = 0; + using var router = new DelegatingFailoverTestRouter( + _ => ++selections == 1 ? failing : working); + + ChatResponse response = + await router.GetStreamingResponseAsync([new(ChatRole.User, "hi")]).ToChatResponseAsync(); + + Assert.Equal("ok", response.Text); + Assert.Equal(2, selections); + } + + [Fact] + public async Task Streaming_FailoverUsesFreshRequestOptionsForEachAttempt() + { + var requestOptions = new ChatOptions + { + Instructions = "caller", + ModelId = "request", + }; + var invokedOptions = new List(); + using var first = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, options, _) => + { + invokedOptions.Add(options!); + options!.ModelId = "changed by first client"; + return ThrowingStream("failed"); + }, + }; + using var second = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, options, _) => + { + invokedOptions.Add(options!); + return YieldUpdates("ok"); + }, + }; + int selections = 0; + using var router = new DelegatingFailoverTestRouter( + _ => ++selections == 1 ? first : second); + + ChatResponse response = await router.GetStreamingResponseAsync( + [new(ChatRole.User, "hi")], + requestOptions).ToChatResponseAsync(); + + Assert.Equal("ok", response.Text); + Assert.Equal(2, selections); + Assert.NotSame(invokedOptions[0], invokedOptions[1]); + Assert.Equal("changed by first client", invokedOptions[0].ModelId); + Assert.Equal("request", invokedOptions[1].ModelId); + Assert.Equal("caller", invokedOptions[1].Instructions); + Assert.Equal("request", requestOptions.ModelId); + } + + [Fact] + public async Task Streaming_CancellationDuringUpdateIsObservedByNextAttempt() + { + using var cancellationSource = new CancellationTokenSource(); + using var inner = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, _, _) => ThrowingStream("failed"), + }; + int selections = 0; + var updates = new List<(FailoverChatClientAttempt attempt, bool isTerminal)>(); + using var router = new DelegatingFailoverTestRouter( + _ => + { + selections++; + return inner; + }, + (_, attempt, isTerminal) => + { + updates.Add((attempt, isTerminal)); + if (!isTerminal) + { + cancellationSource.Cancel(); + } + }); + + await Assert.ThrowsAnyAsync( + () => CollectAsync( + router.GetStreamingResponseAsync( + [new(ChatRole.User, "hi")], + cancellationToken: cancellationSource.Token))); + + Assert.Equal(2, selections); + Assert.Collection( + updates, + update => + { + Assert.NotNull(update.attempt); + Assert.False(update.isTerminal); + }, + update => + { + Assert.NotNull(update.attempt); + Assert.True(update.isTerminal); + }); + } + + [Fact] + public async Task Streaming_CancellationDoesNotSelectNext() + { + using var cancellationSource = new CancellationTokenSource(); + cancellationSource.Cancel(); + using var inner = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, _, _) => ThrowingStream("failed"), + }; + int selections = 0; + FailoverChatClientAttempt? observed = null; + using var router = new DelegatingFailoverTestRouter( + _ => + { + selections++; + return inner; + }, + (_, attempt, isTerminal) => + { + Assert.True(isTerminal); + observed = attempt; + }); + + await Assert.ThrowsAnyAsync( + () => CollectAsync( + router.GetStreamingResponseAsync( + [new(ChatRole.User, "hi")], + cancellationToken: cancellationSource.Token))); + + Assert.Equal(1, selections); + Assert.IsType(observed!.Exception); + Assert.False(observed.OutputCommitted); + Assert.False(observed.ResponseCompleted); + } + + [Fact] + public async Task OperationCanceledException_CanTriggerFailover() + { + bool secondCalled = false; + ChatResponse expected = new(new ChatMessage(ChatRole.Assistant, "ok")); + using var first = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => throw new OperationCanceledException(), + }; + using var second = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => + { + secondCalled = true; + return Task.FromResult(expected); + }, + }; + using var client = new OrderedFailoverChatClient([first, second]); + + ChatResponse response = await client.GetResponseAsync([new(ChatRole.User, "hi")]); + + Assert.Same(expected, response); + Assert.True(secondCalled); + } + + [Fact] + public async Task Streaming_OperationCanceledExceptionCanTriggerFailover() + { + bool secondCalled = false; + using var first = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, _, _) => CanceledStream(CancellationToken.None), + }; + using var second = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, _, _) => + { + secondCalled = true; + return YieldUpdates("ok"); + }, + }; + using var client = new OrderedFailoverChatClient([first, second]); + + ChatResponse response = + await client.GetStreamingResponseAsync([new(ChatRole.User, "hi")]).ToChatResponseAsync(); + + Assert.Equal("ok", response.Text); + Assert.True(secondCalled); + } + + [Fact] + public async Task Streaming_DisposalCancellationCanTriggerFailover() + { + bool secondCalled = false; + var canceledStream = new TrackingAsyncEnumerable( + [], + disposeException: new OperationCanceledException()); + using var first = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, _, _) => canceledStream, + }; + using var second = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, _, _) => + { + secondCalled = true; + return YieldUpdates("ok"); + }, + }; + using var client = new OrderedFailoverChatClient([first, second]); + + ChatResponse response = + await client.GetStreamingResponseAsync([new(ChatRole.User, "hi")]).ToChatResponseAsync(); + + Assert.Equal("ok", response.Text); + Assert.True(secondCalled); + Assert.Equal(1, canceledStream.DisposeCount); + } + + [Fact] + public async Task Streaming_MaximumAttemptsPerRequest_RethrowsLastFailure() + { + int calls = 0; + var terminalUpdates = new List(); + using var inner = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, _, _) => ThrowingStream($"failure {++calls}"), + }; + using var router = new DelegatingFailoverTestRouter( + _ => inner, + (_, attempt, isTerminal) => + { + Assert.NotNull(attempt); + terminalUpdates.Add(isTerminal); + }) + { + MaximumAttemptsPerRequest = 2, + }; + + InvalidOperationException exception = await Assert.ThrowsAsync( + () => CollectAsync(router.GetStreamingResponseAsync([new(ChatRole.User, "hi")]))); + + Assert.Equal("failure 2", exception.Message); + Assert.Equal(2, calls); + Assert.Equal([false, true], terminalUpdates); + } + + [Fact] + public async Task Streaming_StreamCreationFailureFallsBack() + { + using var failing = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, _, _) => throw new InvalidOperationException("failed"), + }; + using var working = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, _, _) => YieldUpdates("ok"), + }; + int selections = 0; + using var router = new DelegatingFailoverTestRouter( + _ => ++selections == 1 ? failing : working); + + ChatResponse response = + await router.GetStreamingResponseAsync([new(ChatRole.User, "hi")]).ToChatResponseAsync(); + + Assert.Equal("ok", response.Text); + Assert.Equal(2, selections); + } + + [Fact] + public async Task Streaming_EnumeratorCreationFailureFallsBack() + { + using var failing = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, _, _) => + new ThrowingGetAsyncEnumeratorEnumerable(new InvalidOperationException("failed")), + }; + using var working = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, _, _) => YieldUpdates("ok"), + }; + int selections = 0; + using var router = new DelegatingFailoverTestRouter( + _ => ++selections == 1 ? failing : working); + + ChatResponse response = + await router.GetStreamingResponseAsync([new(ChatRole.User, "hi")]).ToChatResponseAsync(); + + Assert.Equal("ok", response.Text); + Assert.Equal(2, selections); + } + + [Fact] + public async Task Streaming_CurrentFailureFallsBackBeforeOutput() + { + var currentException = new InvalidOperationException("current failed"); + var failedStream = new ThrowingCurrentAsyncEnumerable(currentException); + using var failing = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, _, _) => failedStream, + }; + using var working = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, _, _) => YieldUpdates("ok"), + }; + FailoverChatClientAttempt? failedAttempt = null; + using var router = new DelegatingFailoverTestRouter( + _ => failedAttempt is null ? failing : working, + (_, attempt, isTerminal) => + { + if (!isTerminal) + { + failedAttempt = attempt; + } + }); + + ChatResponse response = + await router.GetStreamingResponseAsync([new(ChatRole.User, "hi")]).ToChatResponseAsync(); + + Assert.Equal("ok", response.Text); + Assert.Same(currentException, failedAttempt!.Exception); + Assert.False(failedAttempt.OutputCommitted); + Assert.False(failedAttempt.ResponseCompleted); + Assert.Null(failedAttempt.TimeToFirstUpdate); + Assert.Equal(1, failedStream.DisposeCount); + } + + [Fact] + public async Task Streaming_CurrentFailureCancellationIsObservedByNextAttempt() + { + using var cancellationSource = new CancellationTokenSource(); + var failedStream = new ThrowingCurrentAsyncEnumerable(new InvalidOperationException("current failed")); + using var inner = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, _, _) => failedStream, + }; + int selections = 0; + var updates = new List<(FailoverChatClientAttempt attempt, bool isTerminal)>(); + using var router = new DelegatingFailoverTestRouter( + _ => + { + selections++; + return inner; + }, + (_, attempt, isTerminal) => + { + updates.Add((attempt, isTerminal)); + if (!isTerminal) + { + cancellationSource.Cancel(); + } + }); + + await Assert.ThrowsAnyAsync( + () => CollectAsync( + router.GetStreamingResponseAsync( + [new(ChatRole.User, "hi")], + cancellationToken: cancellationSource.Token))); + + Assert.Equal(2, selections); + Assert.Collection( + updates, + update => + { + Assert.NotNull(update.attempt); + Assert.False(update.isTerminal); + }, + update => + { + Assert.NotNull(update.attempt); + Assert.True(update.isTerminal); + }); + Assert.Equal(2, failedStream.DisposeCount); + } + + [Fact] + public async Task Streaming_EmptyStreamDisposalFailureFallsBack() + { + var disposalException = new InvalidOperationException("dispose failed"); + var failedStream = new TrackingAsyncEnumerable([], disposeException: disposalException); + using var failing = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, _, _) => failedStream, + }; + using var working = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, _, _) => YieldUpdates("ok"), + }; + FailoverChatClientAttempt? failedAttempt = null; + FailoverChatClientAttempt? terminalAttempt = null; + using var router = new DelegatingFailoverTestRouter( + _ => failedAttempt is null ? failing : working, + (_, attempt, isTerminal) => + { + if (isTerminal) + { + terminalAttempt = attempt; + } + else + { + failedAttempt = attempt; + } + }); + + ChatResponse response = + await router.GetStreamingResponseAsync([new(ChatRole.User, "hi")]).ToChatResponseAsync(); + + Assert.Equal("ok", response.Text); + Assert.Same(disposalException, failedAttempt!.Exception); + Assert.False(failedAttempt.OutputCommitted); + Assert.False(failedAttempt.ResponseCompleted); + Assert.True(terminalAttempt!.ResponseCompleted); + Assert.Equal(1, failedStream.DisposeCount); + } + + [Fact] + public async Task Streaming_MidStreamFailureIsObservedAndDoesNotReselect() + { + var stream = new TrackingAsyncEnumerable( + [new ChatResponseUpdate(ChatRole.Assistant, "first")], + throwOnMove: 2, + exception: new InvalidOperationException("mid-stream")); + using var inner = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, _, _) => stream, + }; + int selections = 0; + FailoverChatClientAttempt? observed = null; + using var router = new DelegatingFailoverTestRouter( + _ => + { + selections++; + return inner; + }, + (_, attempt, isTerminal) => + { + Assert.True(isTerminal); + observed = attempt; + }); + var updates = new List(); + + async Task ConsumeAsync() + { + await foreach (ChatResponseUpdate update in router.GetStreamingResponseAsync([new(ChatRole.User, "hi")])) + { + updates.Add(update); + } + } + + InvalidOperationException exception = await Assert.ThrowsAsync(ConsumeAsync); + + Assert.Equal("mid-stream", exception.Message); + Assert.Equal("first", Assert.Single(updates).Text); + Assert.Equal(1, selections); + Assert.Same(exception, observed!.Exception); + Assert.True(observed.OutputCommitted); + Assert.False(observed.ResponseCompleted); + Assert.NotNull(observed.TimeToFirstUpdate); + Assert.True(observed.Duration >= observed.TimeToFirstUpdate.Value); + Assert.Equal(1, stream.DisposeCount); + } + + [Fact] + public async Task Streaming_CompletionNotifiesHook() + { + using var inner = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, _, _) => YieldUpdates("a", "b"), + }; + int completions = 0; + FailoverChatClientAttempt? observed = null; + using var router = new DelegatingFailoverTestRouter( + _ => inner, + (_, attempt, isTerminal) => + { + Assert.True(isTerminal); + completions++; + observed = attempt; + }); + + List updates = + await CollectAsync(router.GetStreamingResponseAsync([new(ChatRole.User, "hi")])); + + Assert.Equal(2, updates.Count); + Assert.Equal(1, completions); + Assert.Null(observed!.Exception); + Assert.True(observed.OutputCommitted); + Assert.True(observed.ResponseCompleted); + Assert.NotNull(observed.TimeToFirstUpdate); + } + + [Fact] + public async Task Streaming_EmptyCompletionNotifiesHook() + { + using var inner = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, _, _) => YieldUpdates(), + }; + int updateCount = 0; + using var router = new DelegatingFailoverTestRouter( + _ => inner, + (_, attempt, isTerminal) => + { + updateCount++; + Assert.NotNull(attempt); + Assert.Null(attempt.Exception); + Assert.False(attempt.OutputCommitted); + Assert.True(attempt.ResponseCompleted); + Assert.True(isTerminal); + }); + + List updates = + await CollectAsync(router.GetStreamingResponseAsync([new(ChatRole.User, "hi")])); + + Assert.Empty(updates); + Assert.Equal(1, updateCount); + } + + [Fact] + public async Task Streaming_CallerStopsEarlyNotifiesHookAsIncomplete() + { + using var inner = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, _, _) => YieldUpdates("a", "b"), + }; + int completions = 0; + FailoverChatClientAttempt? observed = null; + using var router = new DelegatingFailoverTestRouter( + _ => inner, + (_, attempt, isTerminal) => + { + Assert.True(isTerminal); + completions++; + observed = attempt; + }); + + await ConsumeOneAsync(router.GetStreamingResponseAsync([new(ChatRole.User, "hi")])); + + Assert.Equal(1, completions); + Assert.Null(observed!.Exception); + Assert.True(observed.OutputCommitted); + Assert.False(observed.ResponseCompleted); + } + + private static async Task> CollectAsync( + IAsyncEnumerable updates) + { + var result = new List(); + await foreach (ChatResponseUpdate update in updates) + { + result.Add(update); + } + + return result; + } + + private static async Task ConsumeOneAsync(IAsyncEnumerable updates) + { + await using IAsyncEnumerator enumerator = updates.GetAsyncEnumerator(); + Assert.True(await enumerator.MoveNextAsync()); + } + + internal static async IAsyncEnumerable YieldUpdates(params string[] texts) + { + foreach (string text in texts) + { + await Task.Yield(); + yield return new ChatResponseUpdate(ChatRole.Assistant, text); + } + } + + internal static async IAsyncEnumerable ThrowingStream(string message) + { + await Task.Yield(); + foreach (int _ in Array.Empty()) + { + yield return new ChatResponseUpdate(ChatRole.Assistant, "never"); + } + + throw new InvalidOperationException(message); + } + + private static async IAsyncEnumerable CanceledStream( + [EnumeratorCancellation] CancellationToken cancellationToken) + { + await Task.Yield(); + foreach (int _ in Array.Empty()) + { + yield return new ChatResponseUpdate(ChatRole.Assistant, "never"); + } + + throw new OperationCanceledException(cancellationToken); + } + + private sealed class DelegatingTestRouter : RoutingChatClient + { + private readonly Func _select; + + public DelegatingTestRouter(Func select) + { + _select = select; + } + + protected override ValueTask SelectClientAsync( + RoutingContext context, + CancellationToken cancellationToken) => + new(new RoutingSelection(_select(context), context.ChatOptions)); + } + + private sealed class DelegatingFailoverTestRouter : FailoverChatClient + { + private readonly Func _select; + private readonly Action? _onRoutingUpdate; + + public DelegatingFailoverTestRouter( + Func select, + Action? onRoutingUpdate = null) + { + _select = select; + _onRoutingUpdate = onRoutingUpdate; + } + + protected override ValueTask SelectClientAsync( + RoutingContext context, + CancellationToken cancellationToken) + { + IChatClient client = _select(context); + return new(client is null ? null! : new RoutingSelection(client, context.ChatOptions)); + } + + protected override ValueTask OnRoutingUpdateAsync( + RoutingContext context, + FailoverChatClientAttempt attempt, + bool isTerminal, + CancellationToken cancellationToken) + { + _onRoutingUpdate?.Invoke(context, attempt, isTerminal); + return default; + } + } + + private sealed class DelegatingSelectionFailoverTestRouter : FailoverChatClient + { + private readonly Func _select; + private readonly Action? _onRoutingUpdate; + + public DelegatingSelectionFailoverTestRouter( + Func select, + Action? onRoutingUpdate = null) + { + _select = select; + _onRoutingUpdate = onRoutingUpdate; + } + + protected override ValueTask SelectClientAsync( + RoutingContext context, + CancellationToken cancellationToken) => + new(_select(context)); + + protected override ValueTask OnRoutingUpdateAsync( + RoutingContext context, + FailoverChatClientAttempt attempt, + bool isTerminal, + CancellationToken cancellationToken) + { + _onRoutingUpdate?.Invoke(context, attempt, isTerminal); + return default; + } + } + + private sealed class ThrowingGetAsyncEnumeratorEnumerable : IAsyncEnumerable + { + private readonly Exception _exception; + + public ThrowingGetAsyncEnumeratorEnumerable(Exception exception) + { + _exception = exception; + } + + public IAsyncEnumerator GetAsyncEnumerator( + CancellationToken cancellationToken = default) => + throw _exception; + } + + private sealed class ThrowingCurrentAsyncEnumerable(Exception exception) : IAsyncEnumerable + { + public int DisposeCount { get; private set; } + + public IAsyncEnumerator GetAsyncEnumerator( + CancellationToken cancellationToken = default) => + new Enumerator(this, exception); + + private sealed class Enumerator( + ThrowingCurrentAsyncEnumerable owner, + Exception exception) : IAsyncEnumerator + { + public ChatResponseUpdate Current => throw exception; + + public ValueTask DisposeAsync() + { + owner.DisposeCount++; + return default; + } + + public ValueTask MoveNextAsync() => new(true); + } + } + + private sealed class TrackingAsyncEnumerable : IAsyncEnumerable + { + private readonly IReadOnlyList _updates; + private readonly int? _throwOnMove; + private readonly Exception? _exception; + private readonly Exception? _disposeException; + + public TrackingAsyncEnumerable( + IReadOnlyList updates, + int? throwOnMove = null, + Exception? exception = null, + Exception? disposeException = null) + { + _updates = updates; + _throwOnMove = throwOnMove; + _exception = exception; + _disposeException = disposeException; + } + + public int DisposeCount { get; private set; } + + public IAsyncEnumerator GetAsyncEnumerator( + CancellationToken cancellationToken = default) => + new Enumerator(this); + + private sealed class Enumerator : IAsyncEnumerator + { + private readonly TrackingAsyncEnumerable _owner; + private int _moveCount; + + public Enumerator(TrackingAsyncEnumerable owner) + { + _owner = owner; + } + + public ChatResponseUpdate Current { get; private set; } = null!; + + public ValueTask MoveNextAsync() + { + _moveCount++; + if (_moveCount == _owner._throwOnMove) + { + throw _owner._exception!; + } + + int index = _moveCount - 1; + if (index >= _owner._updates.Count) + { + return new(false); + } + + Current = _owner._updates[index]; + return new(true); + } + + public ValueTask DisposeAsync() + { + _owner.DisposeCount++; + if (_owner._disposeException is { } exception) + { + throw exception; + } + + return default; + } + } + } +} diff --git a/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/OrderedFailoverChatClientTests.cs b/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/OrderedFailoverChatClientTests.cs new file mode 100644 index 00000000000..5cb2cf34ed8 --- /dev/null +++ b/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/OrderedFailoverChatClientTests.cs @@ -0,0 +1,383 @@ +// 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.Linq; +using System.Threading; +using System.Threading.Tasks; +using Xunit; +using static Microsoft.Extensions.AI.FailoverChatClientTests; + +namespace Microsoft.Extensions.AI; + +public class OrderedFailoverChatClientTests +{ + [Fact] + public void OrderedFailover_RejectsMissingClients() + { + using var inner = new TestChatClient(); + + Assert.Throws(() => new OrderedFailoverChatClient(null!)); + Assert.Throws(() => new OrderedFailoverChatClient([])); + Assert.Throws(() => new OrderedFailoverChatClient([inner, null!])); + Assert.Throws(() => new OrderedFailoverChatClient([inner, inner])); + } + + [Fact] + public async Task OrderedFailover_TriesClientsInOrderAndSkipsFailures() + { + var calls = new List(); + using var first = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => + { + calls.Add("first"); + throw new InvalidOperationException("first failed"); + }, + }; + ChatResponse expected = new(new ChatMessage(ChatRole.Assistant, "ok")); + using var second = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => + { + calls.Add("second"); + return Task.FromResult(expected); + }, + }; + using var client = new OrderedFailoverChatClient([first, second]); + + ChatResponse response = await client.GetResponseAsync([new(ChatRole.User, "hi")]); + + Assert.Same(expected, response); + Assert.Equal(["first", "second"], calls); + } + + [Fact] + public async Task OrderedFailover_DistinguishesValueEqualClients() + { + ChatResponse expected = new(new ChatMessage(ChatRole.Assistant, "ok")); + using var first = new ValueEqualChatClient( + () => throw new InvalidOperationException("failed")); + using var second = new ValueEqualChatClient( + () => Task.FromResult(expected)); + using var client = new OrderedFailoverChatClient([first, second]); + + ChatResponse response = await client.GetResponseAsync([new(ChatRole.User, "hi")]); + + Assert.Same(expected, response); + } + + [Fact] + public async Task OrderedFailover_ExhaustionRethrowsLastFailure() + { + using var first = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => throw new InvalidOperationException("first"), + }; + using var second = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => throw new InvalidOperationException("second"), + }; + using var client = new OrderedFailoverChatClient([first, second]); + + InvalidOperationException exception = await Assert.ThrowsAsync( + () => client.GetResponseAsync([new(ChatRole.User, "hi")])); + + Assert.Equal("second", exception.Message); + } + + [Fact] + public async Task OrderedFailover_StateIsScopedToOneRequest() + { + int firstCalls = 0; + using var first = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => + { + firstCalls++; + throw new InvalidOperationException("failed"); + }, + }; + using var second = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => Task.FromResult(new ChatResponse()), + }; + using var client = new OrderedFailoverChatClient([first, second]); + + _ = await client.GetResponseAsync([new(ChatRole.User, "one")]); + _ = await client.GetResponseAsync([new(ChatRole.User, "two")]); + + Assert.Equal(2, firstCalls); + } + + [Fact] + public async Task OrderedFailover_ConcurrentRequestsHaveIndependentState() + { + const int RequestCount = 10; + int firstCalls = 0; + int secondCalls = 0; + var allStarted = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + using var first = new TestChatClient + { + GetResponseAsyncCallback = async (_, _, _) => + { + if (Interlocked.Increment(ref firstCalls) == RequestCount) + { + allStarted.SetResult(true); + } + + await allStarted.Task; + throw new InvalidOperationException("failed"); + }, + }; + using var second = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => + { + Interlocked.Increment(ref secondCalls); + return Task.FromResult(new ChatResponse()); + }, + }; + using var client = new OrderedFailoverChatClient([first, second]); + + Task[] requests = + [ + .. Enumerable.Range(0, RequestCount).Select(index => + client.GetResponseAsync([new(ChatRole.User, index.ToString())])), + ]; + _ = await Task.WhenAll(requests); + + Assert.Equal(RequestCount, firstCalls); + Assert.Equal(RequestCount, secondCalls); + } + + [Fact] + public async Task OrderedFailover_DoesNotRetainStateWhileStreaming() + { + using var first = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, _, _) => ThrowingStream("failed"), + }; + using var second = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, _, _) => YieldUpdates("first", "second"), + }; + using var client = new OrderedFailoverChatClient([first, second]); + await using IAsyncEnumerator enumerator = + client.GetStreamingResponseAsync([new(ChatRole.User, "hi")]).GetAsyncEnumerator(); + + Assert.True(await enumerator.MoveNextAsync()); + Assert.Equal("first", enumerator.Current.Text); + Assert.Equal(0, GetOrderedFailoverRequestStateCount(client)); + } + + [Fact] + public async Task OrderedFailover_SnapshotsConfiguredClients() + { + ChatResponse expected = new(new ChatMessage(ChatRole.Assistant, "ok")); + using var inner = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => Task.FromResult(expected), + }; + var clients = new List { inner }; + using var failover = new OrderedFailoverChatClient(clients, leaveOpen: true); + clients.Clear(); + + ChatResponse response = await failover.GetResponseAsync([new(ChatRole.User, "hi")]); + + Assert.Same(expected, response); + } + + [Fact] + public async Task OrderedFailover_UsesEachConfiguredClient() + { + var seenOptions = new List<(string? modelId, string? instructions)>(); + using var firstInner = new TestChatClient + { + GetResponseAsyncCallback = (_, options, _) => + { + seenOptions.Add((options?.ModelId, options?.Instructions)); + throw new InvalidOperationException("failed"); + }, + }; + using var first = new ConfigureOptionsChatClient( + firstInner, + options => options.ModelId = "first"); + using var secondInner = new TestChatClient + { + GetResponseAsyncCallback = (_, options, _) => + { + seenOptions.Add((options?.ModelId, options?.Instructions)); + return Task.FromResult(new ChatResponse()); + }, + }; + using var second = new ConfigureOptionsChatClient( + secondInner, + options => options.ModelId = "second"); + using var client = new OrderedFailoverChatClient( + [first, second], + leaveOpen: true); + + _ = await client.GetResponseAsync( + [new(ChatRole.User, "hi")], + new ChatOptions + { + Instructions = "caller", + ModelId = "request", + }); + + Assert.Equal( + [("first", "caller"), ("second", "caller")], + seenOptions); + } + + [Fact] + public async Task OrderedFailover_StreamingFallsBackBeforeFirstUpdate() + { + using var first = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, _, _) => ThrowingStream("failed"), + }; + using var second = new TestChatClient + { + GetStreamingResponseAsyncCallback = (_, _, _) => YieldUpdates("ok"), + }; + using var client = new OrderedFailoverChatClient([first, second]); + + ChatResponse response = + await client.GetStreamingResponseAsync([new(ChatRole.User, "hi")]).ToChatResponseAsync(); + + Assert.Equal("ok", response.Text); + } + + [Fact] + public async Task OrderedFailover_CancellationDoesNotFallback() + { + using var cancellationSource = new CancellationTokenSource(); + cancellationSource.Cancel(); + bool secondCalled = false; + using var first = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => throw new InvalidOperationException("failed"), + }; + using var second = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => + { + secondCalled = true; + return Task.FromResult(new ChatResponse()); + }, + }; + using var client = new OrderedFailoverChatClient([first, second]); + + await Assert.ThrowsAnyAsync( + () => client.GetResponseAsync( + [new(ChatRole.User, "hi")], + cancellationToken: cancellationSource.Token)); + + Assert.False(secondCalled); + } + + [Fact] + public void OrderedFailover_DisposesEachClientOnce() + { +#pragma warning disable CA2000 // Dispose objects before losing scope + var shared = new CountingDisposeClient(); + var other = new CountingDisposeClient(); + var client = new OrderedFailoverChatClient([shared, other]); +#pragma warning restore CA2000 + + client.Dispose(); + client.Dispose(); + + Assert.Equal(1, shared.DisposeCount); + Assert.Equal(1, other.DisposeCount); + } + + [Fact] + public void OrderedFailover_LeaveOpenDoesNotDisposeInnerClients() + { +#pragma warning disable CA2000 // Dispose objects before losing scope + var inner = new CountingDisposeClient(); + var client = new OrderedFailoverChatClient([inner], leaveOpen: true); +#pragma warning restore CA2000 + + client.Dispose(); + + Assert.Equal(0, inner.DisposeCount); + inner.Dispose(); + } + + [Fact] + public async Task NestedRouters_ReturnLeafResponse() + { + ChatResponse expected = new(new ChatMessage(ChatRole.Assistant, "ok")); + using var leaf = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => Task.FromResult(expected), + }; + using var inner = new OrderedFailoverChatClient([leaf]); + using var outer = new OrderedFailoverChatClient([inner]); + + ChatResponse response = await outer.GetResponseAsync([new(ChatRole.User, "hi")]); + + Assert.Same(expected, response); + } + + private static int GetOrderedFailoverRequestStateCount(OrderedFailoverChatClient client) + { + object requestStates = typeof(OrderedFailoverChatClient) + .GetField( + "_requestStates", + System.Reflection.BindingFlags.Instance | System.Reflection.BindingFlags.NonPublic)! + .GetValue(client)!; + return (int)requestStates.GetType().GetProperty("Count")!.GetValue(requestStates)!; + } + + private sealed class CountingDisposeClient : IChatClient + { + public int DisposeCount { get; private set; } + + public Task GetResponseAsync( + IEnumerable messages, ChatOptions? options = null, + CancellationToken cancellationToken = default) => + Task.FromResult(new ChatResponse()); + + public IAsyncEnumerable GetStreamingResponseAsync( + IEnumerable messages, ChatOptions? options = null, + CancellationToken cancellationToken = default) => + throw new NotSupportedException(); + + public object? GetService(Type serviceType, object? serviceKey = null) => null; + + public void Dispose() => DisposeCount++; + + public override bool Equals(object? obj) => obj is CountingDisposeClient; + + public override int GetHashCode() => 0; + } + + private sealed class ValueEqualChatClient(Func> getResponse) : IChatClient + { + public Task GetResponseAsync( + IEnumerable messages, ChatOptions? options = null, + CancellationToken cancellationToken = default) => + getResponse(); + + public IAsyncEnumerable GetStreamingResponseAsync( + IEnumerable messages, ChatOptions? options = null, + CancellationToken cancellationToken = default) => + throw new NotSupportedException(); + + public object? GetService(Type serviceType, object? serviceKey = null) => null; + + public void Dispose() + { + } + + public override bool Equals(object? obj) => obj is ValueEqualChatClient; + + public override int GetHashCode() => 0; + } +} diff --git a/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/SemanticRoutingChatClientTests.cs b/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/SemanticRoutingChatClientTests.cs new file mode 100644 index 00000000000..c41627ae698 --- /dev/null +++ b/test/Libraries/Microsoft.Extensions.AI.Tests/ChatRouting/SemanticRoutingChatClientTests.cs @@ -0,0 +1,285 @@ +// 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.Linq; +using System.Runtime.CompilerServices; +using System.Threading; +using System.Threading.Tasks; +using Xunit; + +namespace Microsoft.Extensions.AI; + +public class SemanticRoutingChatClientTests +{ + [Fact] + public void SemanticRouting_RejectsInvalidConfiguration() + { + using var client = new TestChatClient(); + using var generator = new TestEmbeddingGenerator(); + var profiles = new Dictionary>(ChatClientReferenceComparer.Instance) + { + [client] = ["profile"], + }; + + Assert.Throws(() => + new SemanticRoutingChatClient(null!, profiles, client)); + Assert.Throws(() => + new SemanticRoutingChatClient(generator, null!, client)); + Assert.Throws(() => + new SemanticRoutingChatClient(generator, profiles, null!)); + Assert.Throws(() => + new SemanticRoutingChatClient( + generator, + new Dictionary>(), + client)); + Assert.Throws(() => + new SemanticRoutingChatClient( + generator, + new Dictionary> + { + [client] = [], + }, + client)); + Assert.Throws(() => + new SemanticRoutingChatClient( + generator, + new Dictionary> + { + [client] = [" "], + }, + client)); + Assert.Throws(() => + new SemanticRoutingChatClient(generator, profiles, client, scoreThreshold: 1.1f)); + Assert.Throws(() => + new SemanticRoutingChatClient(generator, profiles, client, topK: 0)); + Assert.Throws(() => + new SemanticRoutingChatClient( + generator, + profiles, + client, + scoreAggregation: (SemanticRoutingChatClient.ScoreAggregation)(-1))); + Assert.Throws(() => + new SemanticRoutingChatClient( + generator, + profiles, + client, + scoreThreshold: 2.1f, + topK: 2, + scoreAggregation: SemanticRoutingChatClient.ScoreAggregation.Sum)); + } + + [Fact] + public async Task SemanticRouting_SelectsBestProfileAndCachesIndex() + { + var vectors = new Dictionary + { + ["code profile"] = [1, 0], + ["writing profile"] = [0, 1], + ["debug this code"] = [1, 0], + }; + int profileBatches = 0; + using var generator = new TestEmbeddingGenerator + { + GenerateAsyncCallback = (values, _, _) => + { + string[] inputs = [.. values]; + if (inputs.Length > 1) + { + profileBatches++; + } + + return Task.FromResult(new GeneratedEmbeddings>( + [.. inputs.Select(input => new Embedding(vectors[input]))])); + }, + }; + ChatResponse expected = new(new ChatMessage(ChatRole.Assistant, "code")); + using var code = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => Task.FromResult(expected), + }; + using var writing = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => + Task.FromResult(new ChatResponse(new ChatMessage(ChatRole.Assistant, "writing"))), + }; + var profiles = new Dictionary>(ChatClientReferenceComparer.Instance) + { + [code] = ["code profile"], + [writing] = ["writing profile"], + }; + using var router = new SemanticRoutingChatClient( + generator, + profiles, + defaultClient: writing, + leaveOpen: true); + + ChatResponse first = await router.GetResponseAsync([new(ChatRole.User, "debug this code")]); + ChatResponse second = await router.GetResponseAsync([new(ChatRole.User, "debug this code")]); + + Assert.Same(expected, first); + Assert.Same(expected, second); + Assert.Equal(1, profileBatches); + } + + [Theory] + [InlineData(SemanticRoutingChatClient.ScoreAggregation.Mean, "code")] + [InlineData(SemanticRoutingChatClient.ScoreAggregation.Sum, "writing")] + public async Task SemanticRouting_AggregatesGlobalTopKByClient( + SemanticRoutingChatClient.ScoreAggregation scoreAggregation, + string expectedResponse) + { + var vectors = new Dictionary + { + ["code"] = [1, 0], + ["writing one"] = [0.8f, 0.6f], + ["writing two"] = [0.8f, -0.6f], + ["query"] = [1, 0], + }; + using var generator = new TestEmbeddingGenerator + { + GenerateAsyncCallback = (values, _, _) => + Task.FromResult(new GeneratedEmbeddings>( + [.. values.Select(input => new Embedding(vectors[input]))])), + }; + using var code = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => + Task.FromResult(new ChatResponse(new ChatMessage(ChatRole.Assistant, "code"))), + }; + using var writing = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => + Task.FromResult(new ChatResponse(new ChatMessage(ChatRole.Assistant, "writing"))), + }; + var profiles = new Dictionary>(ChatClientReferenceComparer.Instance) + { + [code] = ["code"], + [writing] = ["writing one", "writing two"], + }; + using var router = new SemanticRoutingChatClient( + generator, + profiles, + defaultClient: code, + scoreThreshold: scoreAggregation == SemanticRoutingChatClient.ScoreAggregation.Sum ? 1.5f : 0.3f, + topK: 3, + scoreAggregation: scoreAggregation, + leaveOpen: true); + + ChatResponse response = await router.GetResponseAsync([new(ChatRole.User, "query")]); + + Assert.Equal(expectedResponse, response.Text); + } + + [Fact] + public async Task SemanticRouting_UsesDefaultBelowThreshold() + { + var vectors = new Dictionary + { + ["code profile"] = [1, 0], + ["unrelated query"] = [0, 1], + }; + using var generator = new TestEmbeddingGenerator + { + GenerateAsyncCallback = (values, _, _) => + Task.FromResult(new GeneratedEmbeddings>( + [.. values.Select(input => new Embedding(vectors[input]))])), + }; + using var profiled = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => + Task.FromResult(new ChatResponse(new ChatMessage(ChatRole.Assistant, "profiled"))), + }; + ChatResponse expected = new(new ChatMessage(ChatRole.Assistant, "default")); + using var defaultClient = new TestChatClient + { + GetResponseAsyncCallback = (_, _, _) => Task.FromResult(expected), + }; + var profiles = new Dictionary>(ChatClientReferenceComparer.Instance) + { + [profiled] = ["code profile"], + }; + using var router = new SemanticRoutingChatClient( + generator, + profiles, + defaultClient, + scoreThreshold: 0.5f, + leaveOpen: true); + + ChatResponse response = + await router.GetResponseAsync([new(ChatRole.User, "unrelated query")]); + + Assert.Same(expected, response); + } + + [Fact] + public void SemanticRouting_DisposesOwnedResourcesOnce() + { +#pragma warning disable CA2000 // Dispose objects before losing scope + var generator = new CountingEmbeddingGenerator(); + var profiled = new CountingDisposeClient(); + var defaultClient = new CountingDisposeClient(); + var profiles = new Dictionary>(ChatClientReferenceComparer.Instance) + { + [profiled] = ["profile"], + }; + var router = new SemanticRoutingChatClient(generator, profiles, defaultClient); +#pragma warning restore CA2000 + + router.Dispose(); + router.Dispose(); + + Assert.Equal(1, generator.DisposeCount); + Assert.Equal(1, profiled.DisposeCount); + Assert.Equal(1, defaultClient.DisposeCount); + } + + private sealed class CountingDisposeClient : IChatClient + { + public int DisposeCount { get; private set; } + + public Task GetResponseAsync( + IEnumerable messages, ChatOptions? options = null, + CancellationToken cancellationToken = default) => + Task.FromResult(new ChatResponse()); + + public IAsyncEnumerable GetStreamingResponseAsync( + IEnumerable messages, ChatOptions? options = null, + CancellationToken cancellationToken = default) => + throw new NotSupportedException(); + + public object? GetService(Type serviceType, object? serviceKey = null) => null; + + public void Dispose() => DisposeCount++; + + public override bool Equals(object? obj) => obj is CountingDisposeClient; + + public override int GetHashCode() => 0; + } + + private sealed class CountingEmbeddingGenerator : + IEmbeddingGenerator> + { + public int DisposeCount { get; private set; } + + public Task>> GenerateAsync( + IEnumerable values, + EmbeddingGenerationOptions? options = null, + CancellationToken cancellationToken = default) => + throw new NotSupportedException(); + + public object? GetService(Type serviceType, object? serviceKey = null) => null; + + public void Dispose() => DisposeCount++; + } + + private sealed class ChatClientReferenceComparer : IEqualityComparer + { + public static ChatClientReferenceComparer Instance { get; } = new(); + + public bool Equals(IChatClient? x, IChatClient? y) => ReferenceEquals(x, y); + + public int GetHashCode(IChatClient obj) => RuntimeHelpers.GetHashCode(obj); + } +}