|
| 1 | +// Licensed to the .NET Foundation under one or more agreements. |
| 2 | +// The .NET Foundation licenses this file to you under the MIT license. |
| 3 | + |
| 4 | +using System; |
| 5 | +using System.Collections.Generic; |
| 6 | +using System.Diagnostics; |
| 7 | +using System.Threading; |
| 8 | +using System.Threading.Channels; |
| 9 | +using System.Threading.Tasks; |
| 10 | +using Microsoft.Shared.Diagnostics; |
| 11 | + |
| 12 | +#pragma warning disable VSTHRD003 // Avoid awaiting foreign Tasks |
| 13 | + |
| 14 | +namespace Microsoft.Extensions.AI; |
| 15 | + |
| 16 | +/// <summary>A delegating chat client that wraps an inner client with implementations provided by delegates.</summary> |
| 17 | +public sealed class AnonymousDelegatingChatClient : DelegatingChatClient |
| 18 | +{ |
| 19 | + /// <summary>The delegate to use as the implementation of <see cref="CompleteAsync"/>.</summary> |
| 20 | + private readonly Func<IList<ChatMessage>, ChatOptions?, IChatClient, CancellationToken, Task<ChatCompletion>>? _completeFunc; |
| 21 | + |
| 22 | + /// <summary>The delegate to use as the implementation of <see cref="CompleteStreamingAsync"/>.</summary> |
| 23 | + /// <remarks> |
| 24 | + /// When non-<see langword="null"/>, this delegate is used as the implementation of <see cref="CompleteStreamingAsync"/> and |
| 25 | + /// will be invoked with the same arguments as the method itself, along with a reference to the inner client. |
| 26 | + /// When <see langword="null"/>, <see cref="CompleteStreamingAsync"/> will delegate directly to the inner client. |
| 27 | + /// </remarks> |
| 28 | + private readonly Func<IList<ChatMessage>, ChatOptions?, IChatClient, CancellationToken, IAsyncEnumerable<StreamingChatCompletionUpdate>>? _completeStreamingFunc; |
| 29 | + |
| 30 | + /// <summary>The delegate to use as the implementation of both <see cref="CompleteAsync"/> and <see cref="CompleteStreamingAsync"/>.</summary> |
| 31 | + private readonly CompleteSharedFunc? _sharedFunc; |
| 32 | + |
| 33 | + /// <summary> |
| 34 | + /// Initializes a new instance of the <see cref="AnonymousDelegatingChatClient"/> class. |
| 35 | + /// </summary> |
| 36 | + /// <param name="innerClient">The inner client.</param> |
| 37 | + /// <param name="sharedFunc"> |
| 38 | + /// A delegate that provides the implementation for both <see cref="CompleteAsync"/> and <see cref="CompleteStreamingAsync"/>. |
| 39 | + /// In addition to the arguments for the operation, it's provided with a delegate to the inner client that should be |
| 40 | + /// used to perform the operation on the inner client. It will handle both the non-streaming and streaming cases. |
| 41 | + /// </param> |
| 42 | + /// <remarks> |
| 43 | + /// This overload may be used when the anonymous implementation needs to provide pre- and/or post-processing, but doesn't |
| 44 | + /// need to interact with the results of the operation, which will come from the inner client. |
| 45 | + /// </remarks> |
| 46 | + /// <exception cref="ArgumentNullException"><paramref name="innerClient"/> is <see langword="null"/>.</exception> |
| 47 | + /// <exception cref="ArgumentNullException"><paramref name="sharedFunc"/> is <see langword="null"/>.</exception> |
| 48 | + public AnonymousDelegatingChatClient(IChatClient innerClient, CompleteSharedFunc sharedFunc) |
| 49 | + : base(innerClient) |
| 50 | + { |
| 51 | + _ = Throw.IfNull(sharedFunc); |
| 52 | + |
| 53 | + _sharedFunc = sharedFunc; |
| 54 | + } |
| 55 | + |
| 56 | + /// <summary> |
| 57 | + /// Initializes a new instance of the <see cref="AnonymousDelegatingChatClient"/> class. |
| 58 | + /// </summary> |
| 59 | + /// <param name="innerClient">The inner client.</param> |
| 60 | + /// <param name="completeFunc"> |
| 61 | + /// A delegate that provides the implementation for <see cref="CompleteAsync"/>. When <see langword="null"/>, |
| 62 | + /// <paramref name="completeStreamingFunc"/> must be non-null, and the implementation of <see cref="CompleteAsync"/> |
| 63 | + /// will use <paramref name="completeStreamingFunc"/> for the implementation. |
| 64 | + /// </param> |
| 65 | + /// <param name="completeStreamingFunc"> |
| 66 | + /// A delegate that provides the implementation for <see cref="CompleteStreamingAsync"/>. When <see langword="null"/>, |
| 67 | + /// <paramref name="completeFunc"/> must be non-null, and the implementation of <see cref="CompleteStreamingAsync"/> |
| 68 | + /// will use <paramref name="completeFunc"/> for the implementation. |
| 69 | + /// </param> |
| 70 | + /// <exception cref="ArgumentNullException"><paramref name="innerClient"/> is <see langword="null"/>.</exception> |
| 71 | + /// <exception cref="ArgumentNullException">Both <paramref name="completeFunc"/> and <paramref name="completeStreamingFunc"/> are <see langword="null"/>.</exception> |
| 72 | + public AnonymousDelegatingChatClient( |
| 73 | + IChatClient innerClient, |
| 74 | + Func<IList<ChatMessage>, ChatOptions?, IChatClient, CancellationToken, Task<ChatCompletion>>? completeFunc, |
| 75 | + Func<IList<ChatMessage>, ChatOptions?, IChatClient, CancellationToken, IAsyncEnumerable<StreamingChatCompletionUpdate>>? completeStreamingFunc) |
| 76 | + : base(innerClient) |
| 77 | + { |
| 78 | + ThrowIfBothDelegatesNull(completeFunc, completeStreamingFunc); |
| 79 | + |
| 80 | + _completeFunc = completeFunc; |
| 81 | + _completeStreamingFunc = completeStreamingFunc; |
| 82 | + } |
| 83 | + |
| 84 | + /// <inheritdoc/> |
| 85 | + public override Task<ChatCompletion> CompleteAsync( |
| 86 | + IList<ChatMessage> chatMessages, ChatOptions? options = null, CancellationToken cancellationToken = default) |
| 87 | + { |
| 88 | + _ = Throw.IfNull(chatMessages); |
| 89 | + |
| 90 | + if (_sharedFunc is not null) |
| 91 | + { |
| 92 | + return CompleteViaSharedAsync(chatMessages, options, cancellationToken); |
| 93 | + |
| 94 | + async Task<ChatCompletion> CompleteViaSharedAsync(IList<ChatMessage> chatMessages, ChatOptions? options, CancellationToken cancellationToken) |
| 95 | + { |
| 96 | + ChatCompletion? completion = null; |
| 97 | + await _sharedFunc(chatMessages, options, async (chatMessages, options, cancellationToken) => |
| 98 | + { |
| 99 | + completion = await InnerClient.CompleteAsync(chatMessages, options, cancellationToken).ConfigureAwait(false); |
| 100 | + }, cancellationToken).ConfigureAwait(false); |
| 101 | + |
| 102 | + if (completion is null) |
| 103 | + { |
| 104 | + throw new InvalidOperationException("The wrapper completed successfully without producing a ChatCompletion."); |
| 105 | + } |
| 106 | + |
| 107 | + return completion; |
| 108 | + } |
| 109 | + } |
| 110 | + else if (_completeFunc is not null) |
| 111 | + { |
| 112 | + return _completeFunc(chatMessages, options, InnerClient, cancellationToken); |
| 113 | + } |
| 114 | + else |
| 115 | + { |
| 116 | + Debug.Assert(_completeStreamingFunc is not null, "Expected non-null streaming delegate."); |
| 117 | + return _completeStreamingFunc!(chatMessages, options, InnerClient, cancellationToken) |
| 118 | + .ToChatCompletionAsync(coalesceContent: true, cancellationToken); |
| 119 | + } |
| 120 | + } |
| 121 | + |
| 122 | + /// <inheritdoc/> |
| 123 | + public override IAsyncEnumerable<StreamingChatCompletionUpdate> CompleteStreamingAsync( |
| 124 | + IList<ChatMessage> chatMessages, ChatOptions? options = null, CancellationToken cancellationToken = default) |
| 125 | + { |
| 126 | + _ = Throw.IfNull(chatMessages); |
| 127 | + |
| 128 | + if (_sharedFunc is not null) |
| 129 | + { |
| 130 | + var updates = Channel.CreateBounded<StreamingChatCompletionUpdate>(1); |
| 131 | + |
| 132 | +#pragma warning disable CA2016 // explicitly not forwarding the cancellation token, as we need to ensure the channel is always completed |
| 133 | + _ = Task.Run(async () => |
| 134 | +#pragma warning restore CA2016 |
| 135 | + { |
| 136 | + Exception? error = null; |
| 137 | + try |
| 138 | + { |
| 139 | + await _sharedFunc(chatMessages, options, async (chatMessages, options, cancellationToken) => |
| 140 | + { |
| 141 | + await foreach (var update in InnerClient.CompleteStreamingAsync(chatMessages, options, cancellationToken).ConfigureAwait(false)) |
| 142 | + { |
| 143 | + await updates.Writer.WriteAsync(update, cancellationToken).ConfigureAwait(false); |
| 144 | + } |
| 145 | + }, cancellationToken).ConfigureAwait(false); |
| 146 | + } |
| 147 | + catch (Exception ex) |
| 148 | + { |
| 149 | + error = ex; |
| 150 | + throw; |
| 151 | + } |
| 152 | + finally |
| 153 | + { |
| 154 | + _ = updates.Writer.TryComplete(error); |
| 155 | + } |
| 156 | + }); |
| 157 | + |
| 158 | + return updates.Reader.ReadAllAsync(cancellationToken); |
| 159 | + } |
| 160 | + else if (_completeStreamingFunc is not null) |
| 161 | + { |
| 162 | + return _completeStreamingFunc(chatMessages, options, InnerClient, cancellationToken); |
| 163 | + } |
| 164 | + else |
| 165 | + { |
| 166 | + Debug.Assert(_completeFunc is not null, "Expected non-null non-streaming delegate."); |
| 167 | + return CompleteStreamingAsyncViaCompleteAsync(_completeFunc!(chatMessages, options, InnerClient, cancellationToken)); |
| 168 | + |
| 169 | + static async IAsyncEnumerable<StreamingChatCompletionUpdate> CompleteStreamingAsyncViaCompleteAsync(Task<ChatCompletion> task) |
| 170 | + { |
| 171 | + ChatCompletion completion = await task.ConfigureAwait(false); |
| 172 | + foreach (var update in completion.ToStreamingChatCompletionUpdates()) |
| 173 | + { |
| 174 | + yield return update; |
| 175 | + } |
| 176 | + } |
| 177 | + } |
| 178 | + } |
| 179 | + |
| 180 | + /// <summary>Throws an exception if both of the specified delegates are null.</summary> |
| 181 | + /// <exception cref="ArgumentNullException">Both <paramref name="completeFunc"/> and <paramref name="completeStreamingFunc"/> are <see langword="null"/>.</exception> |
| 182 | + internal static void ThrowIfBothDelegatesNull(object? completeFunc, object? completeStreamingFunc) |
| 183 | + { |
| 184 | + if (completeFunc is null && completeStreamingFunc is null) |
| 185 | + { |
| 186 | + Throw.ArgumentNullException(nameof(completeFunc), $"At least one of the {nameof(completeFunc)} or {nameof(completeStreamingFunc)} delegates must be non-null."); |
| 187 | + } |
| 188 | + } |
| 189 | + |
| 190 | + // Design note: |
| 191 | + // The following delegate could juse use Func<...>, but it's defined as a custom delegate type |
| 192 | + // in order to provide better discoverability / documentation / usability around its complicated |
| 193 | + // signature with the nextAsync delegate parameter. |
| 194 | + |
| 195 | + /// <summary> |
| 196 | + /// Represents a method used to call <see cref="IChatClient.CompleteAsync"/> or <see cref="IChatClient.CompleteStreamingAsync"/>. |
| 197 | + /// </summary> |
| 198 | + /// <param name="chatMessages">The chat content to send.</param> |
| 199 | + /// <param name="options">The chat options to configure the request.</param> |
| 200 | + /// <param name="nextAsync"> |
| 201 | + /// A delegate that provides the implementation for the inner client's <see cref="IChatClient.CompleteAsync"/> or |
| 202 | + /// <see cref="IChatClient.CompleteStreamingAsync"/>. It should be invoked to continue the pipeline. It accepts |
| 203 | + /// the chat messages, options, and cancellation token, which are typically the same instances as provided to this method |
| 204 | + /// but need not be. |
| 205 | + /// </param> |
| 206 | + /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param> |
| 207 | + /// <returns>A <see cref="Task"/> that represents the completion of the operation.</returns> |
| 208 | + public delegate Task CompleteSharedFunc( |
| 209 | + IList<ChatMessage> chatMessages, |
| 210 | + ChatOptions? options, |
| 211 | + Func<IList<ChatMessage>, ChatOptions?, CancellationToken, Task> nextAsync, |
| 212 | + CancellationToken cancellationToken); |
| 213 | +} |
0 commit comments