openai/openai-dotnet
Publicmirrored from https://github.com/openai/openai-dotnetAvailable
src/Utility/SseUpdateCollection.cs
203lines · modeblame
58f93c8dShivangiReja2 years ago | 1 | using System; |
9f9f2936Jose Arriaga Maldonado2 years ago | 2 | using System.ClientModel; |
| 3 | using System.ClientModel.Primitives; | |
| 4 | using System.Collections; | |
| 5 | using System.Collections.Generic; | |
674e0f77Stephen Toub2 years ago | 6 | using System.Net.ServerSentEvents; |
9f9f2936Jose Arriaga Maldonado2 years ago | 7 | using System.Text.Json; |
2ab1a942Jose Arriaga Maldonado1 years ago | 8 | using System.Threading; |
9f9f2936Jose Arriaga Maldonado2 years ago | 9 | |
| 10 | #nullable enable | |
| 11 | | |
0ca4c062Jose Arriaga Maldonado1 years ago | 12 | namespace OpenAI; |
9f9f2936Jose Arriaga Maldonado2 years ago | 13 | |
| 14 | /// <summary> | |
0ca4c062Jose Arriaga Maldonado1 years ago | 15 | /// Implementation of collection abstraction over streaming updates. |
9f9f2936Jose Arriaga Maldonado2 years ago | 16 | /// </summary> |
0ca4c062Jose Arriaga Maldonado1 years ago | 17 | internal class SseUpdateCollection<T> : CollectionResult<T> |
9f9f2936Jose Arriaga Maldonado2 years ago | 18 | { |
0ca4c062Jose Arriaga Maldonado1 years ago | 19 | private readonly Func<ClientResult> _sendRequestFunc; |
| 20 | private readonly Func<SseItem<byte[]>, IEnumerable<T>> _eventDeserializerFunc; | |
2ab1a942Jose Arriaga Maldonado1 years ago | 21 | private readonly CancellationToken _cancellationToken; |
9f9f2936Jose Arriaga Maldonado2 years ago | 22 | |
5dce104aJose Arriaga Maldonado1 years ago | 23 | public List<Action> AdditionalDisposalActions { get; } = []; |
| 24 | | |
0ca4c062Jose Arriaga Maldonado1 years ago | 25 | public SseUpdateCollection( |
| 26 | Func<ClientResult> sendRequestFunc, | |
| 27 | Func<JsonElement, ModelReaderWriterOptions, IEnumerable<T>> jsonMultiDeserializerFunc, | |
| 28 | CancellationToken cancellationToken) | |
| 29 | : this( | |
| 30 | sendRequestFunc, | |
| 31 | AsyncSseUpdateCollection<T>.DeserializeSseToMultipleViaJson(jsonMultiDeserializerFunc), | |
| 32 | cancellationToken) | |
| 33 | | |
| 34 | { | |
| 35 | Argument.AssertNotNull(jsonMultiDeserializerFunc, nameof(jsonMultiDeserializerFunc)); | |
| 36 | } | |
| 37 | | |
| 38 | public SseUpdateCollection( | |
| 39 | Func<ClientResult> sendRequestFunc, | |
| 40 | Func<JsonElement, ModelReaderWriterOptions, T> jsonSingleDeserializerFunc, | |
| 41 | CancellationToken cancellationToken) | |
| 42 | : this( | |
| 43 | sendRequestFunc, | |
| 44 | AsyncSseUpdateCollection<T>.DeserializeSseToSingleViaJson(jsonSingleDeserializerFunc), | |
| 45 | cancellationToken) | |
| 46 | { | |
| 47 | Argument.AssertNotNull(jsonSingleDeserializerFunc, nameof(jsonSingleDeserializerFunc)); | |
| 48 | } | |
| 49 | | |
| 50 | public SseUpdateCollection( | |
| 51 | Func<ClientResult> sendRequestFunc, | |
| 52 | Func<SseItem<byte[]>, IEnumerable<T>> eventDeserializerFunc, | |
2ab1a942Jose Arriaga Maldonado1 years ago | 53 | CancellationToken cancellationToken) |
9f9f2936Jose Arriaga Maldonado2 years ago | 54 | { |
0ca4c062Jose Arriaga Maldonado1 years ago | 55 | Argument.AssertNotNull(sendRequestFunc, nameof(sendRequestFunc)); |
| 56 | Argument.AssertNotNull(eventDeserializerFunc, nameof(eventDeserializerFunc)); | |
9f9f2936Jose Arriaga Maldonado2 years ago | 57 | |
0ca4c062Jose Arriaga Maldonado1 years ago | 58 | _sendRequestFunc = sendRequestFunc; |
| 59 | _eventDeserializerFunc = eventDeserializerFunc; | |
2ab1a942Jose Arriaga Maldonado1 years ago | 60 | _cancellationToken = cancellationToken; |
9f9f2936Jose Arriaga Maldonado2 years ago | 61 | } |
| 62 | | |
2ab1a942Jose Arriaga Maldonado1 years ago | 63 | public override ContinuationToken? GetContinuationToken(ClientResult page) |
| 64 | // Continuation is not supported for SSE streams. | |
| 65 | => null; | |
| 66 | | |
| 67 | public override IEnumerable<ClientResult> GetRawPages() | |
9f9f2936Jose Arriaga Maldonado2 years ago | 68 | { |
2ab1a942Jose Arriaga Maldonado1 years ago | 69 | // We don't currently support resuming a dropped connection from the |
| 70 | // last received event, so the response collection has a single element. | |
0ca4c062Jose Arriaga Maldonado1 years ago | 71 | yield return _sendRequestFunc(); |
2ab1a942Jose Arriaga Maldonado1 years ago | 72 | } |
| 73 | | |
0ca4c062Jose Arriaga Maldonado1 years ago | 74 | protected override IEnumerable<T> GetValuesFromPage(ClientResult page) |
2ab1a942Jose Arriaga Maldonado1 years ago | 75 | { |
5dce104aJose Arriaga Maldonado1 years ago | 76 | using IEnumerator<T> enumerator = new SseUpdateEnumerator<T>(_eventDeserializerFunc, page, _cancellationToken, AdditionalDisposalActions); |
2ab1a942Jose Arriaga Maldonado1 years ago | 77 | while (enumerator.MoveNext()) |
| 78 | { | |
| 79 | yield return enumerator.Current; | |
| 80 | } | |
9f9f2936Jose Arriaga Maldonado2 years ago | 81 | } |
| 82 | | |
0ca4c062Jose Arriaga Maldonado1 years ago | 83 | private sealed class SseUpdateEnumerator<U> : IEnumerator<U> |
9f9f2936Jose Arriaga Maldonado2 years ago | 84 | { |
674e0f77Stephen Toub2 years ago | 85 | private static ReadOnlySpan<byte> TerminalData => "[DONE]"u8; |
9f9f2936Jose Arriaga Maldonado2 years ago | 86 | |
5dce104aJose Arriaga Maldonado1 years ago | 87 | private List<Action> _additionalDisposalActions; |
| 88 | | |
2ab1a942Jose Arriaga Maldonado1 years ago | 89 | private readonly CancellationToken _cancellationToken; |
| 90 | private readonly PipelineResponse _response; | |
9f9f2936Jose Arriaga Maldonado2 years ago | 91 | |
| 92 | // These enumerators represent what is effectively a doubly-nested | |
| 93 | // loop over the outer event collection and the inner update collection, | |
| 94 | // i.e.: | |
| 95 | // foreach (var sse in _events) { | |
| 96 | // // get _updates from sse event | |
| 97 | // foreach (var update in _updates) { ... } | |
| 98 | // } | |
674e0f77Stephen Toub2 years ago | 99 | private IEnumerator<SseItem<byte[]>>? _events; |
0ca4c062Jose Arriaga Maldonado1 years ago | 100 | private IEnumerator<U>? _updates; |
| 101 | private readonly Func<SseItem<byte[]>, IEnumerable<U>> _eventDeserializerFunc; | |
9f9f2936Jose Arriaga Maldonado2 years ago | 102 | |
0ca4c062Jose Arriaga Maldonado1 years ago | 103 | private U? _current; |
9f9f2936Jose Arriaga Maldonado2 years ago | 104 | private bool _started; |
| 105 | | |
0ca4c062Jose Arriaga Maldonado1 years ago | 106 | public SseUpdateEnumerator( |
| 107 | Func<SseItem<byte[]>, IEnumerable<U>> eventDeserializerFunc, | |
| 108 | ClientResult page, | |
5dce104aJose Arriaga Maldonado1 years ago | 109 | CancellationToken cancellationToken, |
| 110 | List<Action> additionalDisposalActions) | |
9f9f2936Jose Arriaga Maldonado2 years ago | 111 | { |
0ca4c062Jose Arriaga Maldonado1 years ago | 112 | Argument.AssertNotNull(eventDeserializerFunc, nameof(eventDeserializerFunc)); |
2ab1a942Jose Arriaga Maldonado1 years ago | 113 | Argument.AssertNotNull(page, nameof(page)); |
9f9f2936Jose Arriaga Maldonado2 years ago | 114 | |
0ca4c062Jose Arriaga Maldonado1 years ago | 115 | _eventDeserializerFunc = eventDeserializerFunc; |
2ab1a942Jose Arriaga Maldonado1 years ago | 116 | _response = page.GetRawResponse(); |
| 117 | _cancellationToken = cancellationToken; | |
5dce104aJose Arriaga Maldonado1 years ago | 118 | _additionalDisposalActions = additionalDisposalActions; |
9f9f2936Jose Arriaga Maldonado2 years ago | 119 | } |
| 120 | | |
0ca4c062Jose Arriaga Maldonado1 years ago | 121 | U IEnumerator<U>.Current => _current!; |
9f9f2936Jose Arriaga Maldonado2 years ago | 122 | |
2ab1a942Jose Arriaga Maldonado1 years ago | 123 | object IEnumerator.Current => _current!; |
9f9f2936Jose Arriaga Maldonado2 years ago | 124 | |
| 125 | public bool MoveNext() | |
| 126 | { | |
| 127 | if (_events is null && _started) | |
| 128 | { | |
0ca4c062Jose Arriaga Maldonado1 years ago | 129 | throw new ObjectDisposedException(typeof(U).Name); |
9f9f2936Jose Arriaga Maldonado2 years ago | 130 | } |
| 131 | | |
2ab1a942Jose Arriaga Maldonado1 years ago | 132 | _cancellationToken.ThrowIfCancellationRequested(); |
9f9f2936Jose Arriaga Maldonado2 years ago | 133 | _events ??= CreateEventEnumerator(); |
| 134 | _started = true; | |
| 135 | | |
| 136 | if (_updates is not null && _updates.MoveNext()) | |
| 137 | { | |
| 138 | _current = _updates.Current; | |
| 139 | return true; | |
| 140 | } | |
| 141 | | |
| 142 | if (_events.MoveNext()) | |
| 143 | { | |
674e0f77Stephen Toub2 years ago | 144 | if (_events.Current.Data.AsSpan().SequenceEqual(TerminalData)) |
9f9f2936Jose Arriaga Maldonado2 years ago | 145 | { |
| 146 | _current = default; | |
| 147 | return false; | |
| 148 | } | |
| 149 | | |
0ca4c062Jose Arriaga Maldonado1 years ago | 150 | _updates = _eventDeserializerFunc.Invoke(_events.Current).GetEnumerator(); |
9f9f2936Jose Arriaga Maldonado2 years ago | 151 | |
| 152 | if (_updates.MoveNext()) | |
| 153 | { | |
| 154 | _current = _updates.Current; | |
| 155 | return true; | |
| 156 | } | |
| 157 | } | |
| 158 | | |
| 159 | _current = default; | |
| 160 | return false; | |
| 161 | } | |
| 162 | | |
674e0f77Stephen Toub2 years ago | 163 | private IEnumerator<SseItem<byte[]>> CreateEventEnumerator() |
9f9f2936Jose Arriaga Maldonado2 years ago | 164 | { |
2ab1a942Jose Arriaga Maldonado1 years ago | 165 | if (_response.ContentStream is null) |
9f9f2936Jose Arriaga Maldonado2 years ago | 166 | { |
| 167 | throw new InvalidOperationException("Unable to create result from response with null ContentStream"); | |
| 168 | } | |
| 169 | | |
2ab1a942Jose Arriaga Maldonado1 years ago | 170 | IEnumerable<SseItem<byte[]>> enumerable = SseParser.Create(_response.ContentStream, (_, bytes) => bytes.ToArray()).Enumerate(); |
9f9f2936Jose Arriaga Maldonado2 years ago | 171 | return enumerable.GetEnumerator(); |
| 172 | } | |
| 173 | | |
| 174 | public void Reset() | |
| 175 | { | |
| 176 | throw new NotSupportedException("Cannot seek back in an SSE stream."); | |
| 177 | } | |
| 178 | | |
| 179 | public void Dispose() | |
| 180 | { | |
| 181 | Dispose(true); | |
| 182 | GC.SuppressFinalize(this); | |
| 183 | } | |
| 184 | | |
| 185 | private void Dispose(bool disposing) | |
| 186 | { | |
| 187 | if (disposing && _events is not null) | |
| 188 | { | |
| 189 | _events.Dispose(); | |
| 190 | _events = null; | |
| 191 | | |
2ab1a942Jose Arriaga Maldonado1 years ago | 192 | // Dispose the response so we don't leave the network connection open. |
| 193 | _response?.Dispose(); | |
9f9f2936Jose Arriaga Maldonado2 years ago | 194 | } |
5dce104aJose Arriaga Maldonado1 years ago | 195 | |
| 196 | foreach (Action additionalDisposalAction in _additionalDisposalActions ?? []) | |
| 197 | { | |
| 198 | additionalDisposalAction?.Invoke(); | |
| 199 | } | |
| 200 | _additionalDisposalActions?.Clear(); | |
9f9f2936Jose Arriaga Maldonado2 years ago | 201 | } |
| 202 | } | |
| 203 | } |