GitVita
New repositoryImport repositoryNew issue
Your repositories Your dashboard Sign in

openai/openai-dotnet

Public

mirrored from https://github.com/openai/openai-dotnetAvailable

Watch0Fork0Star0
CodeCommitsIssuesPull requestsActionsInsightsSecurity
fe4d6cf1ccfea48fa5c2baf67103495ac9b459dd

Branches

Tags

  • No tags available.
0Branches0Tags
Go to file
Add file
Code

Clone

HTTPS

Download ZIP

openai-dotnet/OpenAI/src/Utility

OpenAI/src/Utility/AsyncSseUpdateCollection.cs

255lines · modecode

RawDownload
Latest commit unavailable.
1using System;
2using System.ClientModel;
3using System.ClientModel.Primitives;
4using System.Collections.Generic;
5using System.Net.ServerSentEvents;
6using System.Runtime.CompilerServices;
7using System.Text.Json;
8using System.Threading;
9using System.Threading.Tasks;
10
11#nullable enable
12
13namespace OpenAI;
14
15/// <summary>
16/// Implementation of collection abstraction over streaming chat updates.
17/// </summary>
18internal class AsyncSseUpdateCollection<T> : AsyncCollectionResult<T>
19{
20 private readonly Func<Task<ClientResult>> _sendRequestAsync;
21 private readonly Func<SseItem<byte[]>, IEnumerable<T>> _eventDeserializerFunc;
22 private readonly CancellationToken _cancellationToken;
23
24 public List<Action> AdditionalDisposalActions { get; } = [];
25
26 public AsyncSseUpdateCollection(
27 Func<Task<ClientResult>> sendRequestAsync,
28 Func<JsonElement, ModelReaderWriterOptions, IEnumerable<T>> jsonMultiDeserializerFunc,
29 CancellationToken cancellationToken)
30 : this(
31 sendRequestAsync,
32 DeserializeSseToMultipleViaJson(jsonMultiDeserializerFunc),
33 cancellationToken)
34 {
35 Argument.AssertNotNull(jsonMultiDeserializerFunc, nameof(jsonMultiDeserializerFunc));
36 }
37
38 public AsyncSseUpdateCollection(
39 Func<Task<ClientResult>> sendRequestAsync,
40 Func<JsonElement, ModelReaderWriterOptions, T> jsonSingleDeserializerFunc,
41 CancellationToken cancellationToken)
42 : this(
43 sendRequestAsync,
44 DeserializeSseToSingleViaJson(jsonSingleDeserializerFunc),
45 cancellationToken)
46 {
47 Argument.AssertNotNull(jsonSingleDeserializerFunc, nameof(jsonSingleDeserializerFunc));
48 }
49
50 public AsyncSseUpdateCollection(
51 Func<Task<ClientResult>> sendRequestAsync,
52 Func<JsonElement, BinaryData, ModelReaderWriterOptions, IEnumerable<T>> jsonMultiDeserializerFunc,
53 CancellationToken cancellationToken)
54 : this(
55 sendRequestAsync,
56 DeserializeSseToMultipleViaJson(jsonMultiDeserializerFunc),
57 cancellationToken)
58 {
59 Argument.AssertNotNull(jsonMultiDeserializerFunc, nameof(jsonMultiDeserializerFunc));
60 }
61
62 public AsyncSseUpdateCollection(
63 Func<Task<ClientResult>> sendRequestAsync,
64 Func<JsonElement, BinaryData, ModelReaderWriterOptions, T> jsonSingleDeserializerFunc,
65 CancellationToken cancellationToken)
66 : this(
67 sendRequestAsync,
68 DeserializeSseToSingleViaJson(jsonSingleDeserializerFunc),
69 cancellationToken)
70 {
71 Argument.AssertNotNull(jsonSingleDeserializerFunc, nameof(jsonSingleDeserializerFunc));
72 }
73
74 public AsyncSseUpdateCollection(
75 Func<Task<ClientResult>> sendRequestAsync,
76 Func<SseItem<byte[]>, IEnumerable<T>> eventDeserializerFunc,
77 CancellationToken cancellationToken)
78 {
79 Argument.AssertNotNull(sendRequestAsync, nameof(sendRequestAsync));
80 Argument.AssertNotNull(eventDeserializerFunc, nameof(eventDeserializerFunc));
81
82 _sendRequestAsync = sendRequestAsync;
83 _eventDeserializerFunc = eventDeserializerFunc;
84 _cancellationToken = cancellationToken;
85 }
86
87 public override ContinuationToken? GetContinuationToken(ClientResult page)
88 // Continuation is not supported for SSE streams.
89 => null;
90
91 public async override IAsyncEnumerable<ClientResult> GetRawPagesAsync()
92 {
93 // We don't currently support resuming a dropped connection from the
94 // last received event, so the response collection has a single element.
95 yield return await _sendRequestAsync();
96 }
97
98 protected async override IAsyncEnumerable<T> GetValuesFromPageAsync(ClientResult page)
99 {
100 await using IAsyncEnumerator<T> enumerator = new AsyncSseUpdateEnumerator<T>(_eventDeserializerFunc, page, _cancellationToken, AdditionalDisposalActions);
101
102 while (await enumerator.MoveNextAsync().ConfigureAwait(false))
103 {
104 yield return enumerator.Current;
105 }
106 }
107
108 [MethodImpl(MethodImplOptions.AggressiveInlining)]
109 internal static Func<SseItem<byte[]>, IEnumerable<U>> DeserializeSseToMultipleViaJson<U>(
110 Func<JsonElement, ModelReaderWriterOptions, IEnumerable<U>> jsonDeserializationFunc)
111 {
112 return (item) =>
113 {
114 using JsonDocument document = JsonDocument.Parse(item.Data);
115 return jsonDeserializationFunc.Invoke(document.RootElement, ModelSerializationExtensions.WireOptions);
116 };
117 }
118
119 [MethodImpl(MethodImplOptions.AggressiveInlining)]
120 internal static Func<SseItem<byte[]>, IEnumerable<U>> DeserializeSseToSingleViaJson<U>(
121 Func<JsonElement, ModelReaderWriterOptions, U> jsonSingleDeserializationFunc)
122 => DeserializeSseToMultipleViaJson<U>((e, o) => [jsonSingleDeserializationFunc.Invoke(e, o)]);
123
124 [MethodImpl(MethodImplOptions.AggressiveInlining)]
125 internal static Func<SseItem<byte[]>, IEnumerable<U>> DeserializeSseToMultipleViaJson<U>(
126 Func<JsonElement, BinaryData, ModelReaderWriterOptions, IEnumerable<U>> jsonDeserializationFunc)
127 {
128 return (item) =>
129 {
130 using JsonDocument document = JsonDocument.Parse(item.Data);
131 return jsonDeserializationFunc.Invoke(document.RootElement, BinaryData.FromBytes(item.Data), ModelSerializationExtensions.WireOptions);
132 };
133 }
134
135 [MethodImpl(MethodImplOptions.AggressiveInlining)]
136 internal static Func<SseItem<byte[]>, IEnumerable<U>> DeserializeSseToSingleViaJson<U>(
137 Func<JsonElement, BinaryData, ModelReaderWriterOptions, U> jsonSingleDeserializationFunc)
138 => DeserializeSseToMultipleViaJson<U>((e, d, o) => [jsonSingleDeserializationFunc.Invoke(e, d, o)]);
139
140 private sealed class AsyncSseUpdateEnumerator<U> : IAsyncEnumerator<U>
141 {
142 private static ReadOnlySpan<byte> TerminalData => "[DONE]"u8;
143
144 private List<Action> _additionalDisposalActions;
145
146 private readonly CancellationToken _cancellationToken;
147 private readonly PipelineResponse _response;
148
149 // These enumerators represent what is effectively a doubly-nested
150 // loop over the outer event collection and the inner update collection,
151 // i.e.:
152 // foreach (var sse in _events) {
153 // // get _updates from sse event
154 // foreach (var update in _updates) { ... }
155 // }
156 private IAsyncEnumerator<SseItem<byte[]>>? _events;
157 private IEnumerator<U>? _updates;
158 private readonly Func<SseItem<byte[]>, IEnumerable<U>> _deserializerFunc;
159
160 private U? _current;
161 private bool _started;
162
163 public AsyncSseUpdateEnumerator(
164 Func<SseItem<byte[]>, IEnumerable<U>> deserializerFunc,
165 ClientResult page,
166 CancellationToken cancellationToken,
167 List<Action> additionalDisposalActions)
168 {
169 Argument.AssertNotNull(page, nameof(page));
170
171 _deserializerFunc = deserializerFunc;
172 _response = page.GetRawResponse();
173 _cancellationToken = cancellationToken;
174 _additionalDisposalActions = additionalDisposalActions;
175 }
176
177 U IAsyncEnumerator<U>.Current => _current!;
178
179 async ValueTask<bool> IAsyncEnumerator<U>.MoveNextAsync()
180 {
181 if (_events is null && _started)
182 {
183 throw new ObjectDisposedException(nameof(AsyncSseUpdateEnumerator<U>));
184 }
185
186 _cancellationToken.ThrowIfCancellationRequested();
187 _events ??= CreateEventEnumeratorAsync();
188 _started = true;
189
190 if (_updates is not null && _updates.MoveNext())
191 {
192 _current = _updates.Current;
193 return true;
194 }
195
196 if (await _events.MoveNextAsync().ConfigureAwait(false))
197 {
198 if (_events.Current.Data.AsSpan().SequenceEqual(TerminalData))
199 {
200 _current = default;
201 return false;
202 }
203
204 _updates = _deserializerFunc
205 .Invoke(_events.Current)
206 .GetEnumerator();
207
208 if (_updates.MoveNext())
209 {
210 _current = _updates.Current;
211 return true;
212 }
213 }
214
215 _current = default;
216 return false;
217 }
218
219 private IAsyncEnumerator<SseItem<byte[]>> CreateEventEnumeratorAsync()
220 {
221 if (_response.ContentStream is null)
222 {
223 throw new InvalidOperationException("Unable to create result from response with null ContentStream");
224 }
225
226 IAsyncEnumerable<SseItem<byte[]>> enumerable = SseParser.Create(_response.ContentStream, (_, bytes) => bytes.ToArray()).EnumerateAsync();
227 return enumerable.GetAsyncEnumerator(_cancellationToken);
228 }
229
230 public async ValueTask DisposeAsync()
231 {
232 await DisposeAsyncCore().ConfigureAwait(false);
233
234 GC.SuppressFinalize(this);
235 }
236
237 private async ValueTask DisposeAsyncCore()
238 {
239 if (_events is not null)
240 {
241 await _events.DisposeAsync().ConfigureAwait(false);
242 _events = null;
243
244 // Dispose the response so we don't leave the network connection open.
245 _response?.Dispose();
246 }
247
248 foreach (Action additionalDisposalAction in _additionalDisposalActions ?? [])
249 {
250 additionalDisposalAction.Invoke();
251 }
252 _additionalDisposalActions?.Clear();
253 }
254 }
255}
256