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
OpenAI_2.2.0

Branches

Tags

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

Clone

HTTPS

Download ZIP

openai-dotnet/src/Utility

src/Utility/AsyncSseUpdateCollection.cs

215lines · 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<SseItem<byte[]>, IEnumerable<T>> eventDeserializerFunc,
53 CancellationToken cancellationToken)
54 {
55 Argument.AssertNotNull(sendRequestAsync, nameof(sendRequestAsync));
56 Argument.AssertNotNull(eventDeserializerFunc, nameof(eventDeserializerFunc));
57
58 _sendRequestAsync = sendRequestAsync;
59 _eventDeserializerFunc = eventDeserializerFunc;
60 _cancellationToken = cancellationToken;
61 }
62
63 public override ContinuationToken? GetContinuationToken(ClientResult page)
64 // Continuation is not supported for SSE streams.
65 => null;
66
67 public async override IAsyncEnumerable<ClientResult> GetRawPagesAsync()
68 {
69 // We don't currently support resuming a dropped connection from the
70 // last received event, so the response collection has a single element.
71 yield return await _sendRequestAsync();
72 }
73
74 protected async override IAsyncEnumerable<T> GetValuesFromPageAsync(ClientResult page)
75 {
76 await using IAsyncEnumerator<T> enumerator = new AsyncSseUpdateEnumerator<T>(_eventDeserializerFunc, page, _cancellationToken, AdditionalDisposalActions);
77
78 while (await enumerator.MoveNextAsync().ConfigureAwait(false))
79 {
80 yield return enumerator.Current;
81 }
82 }
83
84 [MethodImpl(MethodImplOptions.AggressiveInlining)]
85 internal static Func<SseItem<byte[]>, IEnumerable<U>> DeserializeSseToMultipleViaJson<U>(
86 Func<JsonElement, ModelReaderWriterOptions, IEnumerable<U>> jsonDeserializationFunc)
87 {
88 return (item) =>
89 {
90 using JsonDocument document = JsonDocument.Parse(item.Data);
91 return jsonDeserializationFunc.Invoke(document.RootElement, ModelSerializationExtensions.WireOptions);
92 };
93 }
94
95 [MethodImpl(MethodImplOptions.AggressiveInlining)]
96 internal static Func<SseItem<byte[]>, IEnumerable<U>> DeserializeSseToSingleViaJson<U>(
97 Func<JsonElement, ModelReaderWriterOptions, U> jsonSingleDeserializationFunc)
98 => DeserializeSseToMultipleViaJson<U>((e, o) => [jsonSingleDeserializationFunc.Invoke(e, o)]);
99
100 private sealed class AsyncSseUpdateEnumerator<U> : IAsyncEnumerator<U>
101 {
102 private static ReadOnlySpan<byte> TerminalData => "[DONE]"u8;
103
104 private List<Action> _additionalDisposalActions;
105
106 private readonly CancellationToken _cancellationToken;
107 private readonly PipelineResponse _response;
108
109 // These enumerators represent what is effectively a doubly-nested
110 // loop over the outer event collection and the inner update collection,
111 // i.e.:
112 // foreach (var sse in _events) {
113 // // get _updates from sse event
114 // foreach (var update in _updates) { ... }
115 // }
116 private IAsyncEnumerator<SseItem<byte[]>>? _events;
117 private IEnumerator<U>? _updates;
118 private readonly Func<SseItem<byte[]>, IEnumerable<U>> _deserializerFunc;
119
120 private U? _current;
121 private bool _started;
122
123 public AsyncSseUpdateEnumerator(
124 Func<SseItem<byte[]>, IEnumerable<U>> deserializerFunc,
125 ClientResult page,
126 CancellationToken cancellationToken,
127 List<Action> additionalDisposalActions)
128 {
129 Argument.AssertNotNull(page, nameof(page));
130
131 _deserializerFunc = deserializerFunc;
132 _response = page.GetRawResponse();
133 _cancellationToken = cancellationToken;
134 _additionalDisposalActions = additionalDisposalActions;
135 }
136
137 U IAsyncEnumerator<U>.Current => _current!;
138
139 async ValueTask<bool> IAsyncEnumerator<U>.MoveNextAsync()
140 {
141 if (_events is null && _started)
142 {
143 throw new ObjectDisposedException(nameof(AsyncSseUpdateEnumerator<U>));
144 }
145
146 _cancellationToken.ThrowIfCancellationRequested();
147 _events ??= CreateEventEnumeratorAsync();
148 _started = true;
149
150 if (_updates is not null && _updates.MoveNext())
151 {
152 _current = _updates.Current;
153 return true;
154 }
155
156 if (await _events.MoveNextAsync().ConfigureAwait(false))
157 {
158 if (_events.Current.Data.AsSpan().SequenceEqual(TerminalData))
159 {
160 _current = default;
161 return false;
162 }
163
164 _updates = _deserializerFunc
165 .Invoke(_events.Current)
166 .GetEnumerator();
167
168 if (_updates.MoveNext())
169 {
170 _current = _updates.Current;
171 return true;
172 }
173 }
174
175 _current = default;
176 return false;
177 }
178
179 private IAsyncEnumerator<SseItem<byte[]>> CreateEventEnumeratorAsync()
180 {
181 if (_response.ContentStream is null)
182 {
183 throw new InvalidOperationException("Unable to create result from response with null ContentStream");
184 }
185
186 IAsyncEnumerable<SseItem<byte[]>> enumerable = SseParser.Create(_response.ContentStream, (_, bytes) => bytes.ToArray()).EnumerateAsync();
187 return enumerable.GetAsyncEnumerator(_cancellationToken);
188 }
189
190 public async ValueTask DisposeAsync()
191 {
192 await DisposeAsyncCore().ConfigureAwait(false);
193
194 GC.SuppressFinalize(this);
195 }
196
197 private async ValueTask DisposeAsyncCore()
198 {
199 if (_events is not null)
200 {
201 await _events.DisposeAsync().ConfigureAwait(false);
202 _events = null;
203
204 // Dispose the response so we don't leave the network connection open.
205 _response?.Dispose();
206 }
207
208 foreach (Action additionalDisposalAction in _additionalDisposalActions ?? [])
209 {
210 additionalDisposalAction.Invoke();
211 }
212 _additionalDisposalActions?.Clear();
213 }
214 }
215}
216