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
jsquire-patch-1

Branches

Tags

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

Clone

HTTPS

Download ZIP

openai-dotnet/src/Utility

src/Utility/AsyncSseUpdateCollection.cs

202lines · modeblame

RawDownload
Latest commit unavailable.
58f93c8dShivangiReja2 years ago1using System;
9f9f2936Jose Arriaga Maldonado2 years ago2using System.ClientModel;
3using System.ClientModel.Primitives;
4using System.Collections.Generic;
674e0f77Stephen Toub2 years ago5using System.Net.ServerSentEvents;
0ca4c062Jose Arriaga Maldonado1 years ago6using System.Runtime.CompilerServices;
9f9f2936Jose Arriaga Maldonado2 years ago7using System.Text.Json;
8using System.Threading;
9using System.Threading.Tasks;
10
11#nullable enable
12
0ca4c062Jose Arriaga Maldonado1 years ago13namespace OpenAI;
9f9f2936Jose Arriaga Maldonado2 years ago14
15/// <summary>
16/// Implementation of collection abstraction over streaming chat updates.
17/// </summary>
0ca4c062Jose Arriaga Maldonado1 years ago18internal class AsyncSseUpdateCollection<T> : AsyncCollectionResult<T>
9f9f2936Jose Arriaga Maldonado2 years ago19{
2ab1a942Jose Arriaga Maldonado1 years ago20private readonly Func<Task<ClientResult>> _sendRequestAsync;
0ca4c062Jose Arriaga Maldonado1 years ago21private readonly Func<SseItem<byte[]>, IEnumerable<T>> _eventDeserializerFunc;
2ab1a942Jose Arriaga Maldonado1 years ago22private readonly CancellationToken _cancellationToken;
9f9f2936Jose Arriaga Maldonado2 years ago23
0ca4c062Jose Arriaga Maldonado1 years ago24public AsyncSseUpdateCollection(
2ab1a942Jose Arriaga Maldonado1 years ago25Func<Task<ClientResult>> sendRequestAsync,
0ca4c062Jose Arriaga Maldonado1 years ago26Func<JsonElement, ModelReaderWriterOptions, IEnumerable<T>> jsonMultiDeserializerFunc,
27CancellationToken cancellationToken)
28: this(
29sendRequestAsync,
30DeserializeSseToMultipleViaJson(jsonMultiDeserializerFunc),
31cancellationToken)
32{
33Argument.AssertNotNull(jsonMultiDeserializerFunc, nameof(jsonMultiDeserializerFunc));
34}
35
36public AsyncSseUpdateCollection(
37Func<Task<ClientResult>> sendRequestAsync,
38Func<JsonElement, ModelReaderWriterOptions, T> jsonSingleDeserializerFunc,
39CancellationToken cancellationToken)
40: this(
41sendRequestAsync,
42DeserializeSseToSingleViaJson(jsonSingleDeserializerFunc),
43cancellationToken)
44{
45Argument.AssertNotNull(jsonSingleDeserializerFunc, nameof(jsonSingleDeserializerFunc));
46}
47
48public AsyncSseUpdateCollection(
49Func<Task<ClientResult>> sendRequestAsync,
50Func<SseItem<byte[]>, IEnumerable<T>> eventDeserializerFunc,
2ab1a942Jose Arriaga Maldonado1 years ago51CancellationToken cancellationToken)
9f9f2936Jose Arriaga Maldonado2 years ago52{
2ab1a942Jose Arriaga Maldonado1 years ago53Argument.AssertNotNull(sendRequestAsync, nameof(sendRequestAsync));
0ca4c062Jose Arriaga Maldonado1 years ago54Argument.AssertNotNull(eventDeserializerFunc, nameof(eventDeserializerFunc));
9f9f2936Jose Arriaga Maldonado2 years ago55
2ab1a942Jose Arriaga Maldonado1 years ago56_sendRequestAsync = sendRequestAsync;
0ca4c062Jose Arriaga Maldonado1 years ago57_eventDeserializerFunc = eventDeserializerFunc;
2ab1a942Jose Arriaga Maldonado1 years ago58_cancellationToken = cancellationToken;
9f9f2936Jose Arriaga Maldonado2 years ago59}
60
2ab1a942Jose Arriaga Maldonado1 years ago61public override ContinuationToken? GetContinuationToken(ClientResult page)
62// Continuation is not supported for SSE streams.
63=> null;
64
65public async override IAsyncEnumerable<ClientResult> GetRawPagesAsync()
66{
67// We don't currently support resuming a dropped connection from the
68// last received event, so the response collection has a single element.
69yield return await _sendRequestAsync();
70}
71
0ca4c062Jose Arriaga Maldonado1 years ago72protected async override IAsyncEnumerable<T> GetValuesFromPageAsync(ClientResult page)
9f9f2936Jose Arriaga Maldonado2 years ago73{
0ca4c062Jose Arriaga Maldonado1 years ago74await using IAsyncEnumerator<T> enumerator = new AsyncSseUpdateEnumerator<T>(_eventDeserializerFunc, page, _cancellationToken);
2ab1a942Jose Arriaga Maldonado1 years ago75while (await enumerator.MoveNextAsync().ConfigureAwait(false))
76{
77yield return enumerator.Current;
78}
9f9f2936Jose Arriaga Maldonado2 years ago79}
80
0ca4c062Jose Arriaga Maldonado1 years ago81[MethodImpl(MethodImplOptions.AggressiveInlining)]
82internal static Func<SseItem<byte[]>, IEnumerable<U>> DeserializeSseToMultipleViaJson<U>(
83Func<JsonElement, ModelReaderWriterOptions, IEnumerable<U>> jsonDeserializationFunc)
84{
85return (item) =>
86{
87using JsonDocument document = JsonDocument.Parse(item.Data);
88return jsonDeserializationFunc.Invoke(document.RootElement, ModelSerializationExtensions.WireOptions);
89};
90}
91
92[MethodImpl(MethodImplOptions.AggressiveInlining)]
93internal static Func<SseItem<byte[]>, IEnumerable<U>> DeserializeSseToSingleViaJson<U>(
94Func<JsonElement, ModelReaderWriterOptions, U> jsonSingleDeserializationFunc)
95=> DeserializeSseToMultipleViaJson<U>((e, o) => [jsonSingleDeserializationFunc.Invoke(e, o)]);
96
97private sealed class AsyncSseUpdateEnumerator<U> : IAsyncEnumerator<U>
9f9f2936Jose Arriaga Maldonado2 years ago98{
674e0f77Stephen Toub2 years ago99private static ReadOnlySpan<byte> TerminalData => "[DONE]"u8;
9f9f2936Jose Arriaga Maldonado2 years ago100
101private readonly CancellationToken _cancellationToken;
2ab1a942Jose Arriaga Maldonado1 years ago102private readonly PipelineResponse _response;
9f9f2936Jose Arriaga Maldonado2 years ago103
104// These enumerators represent what is effectively a doubly-nested
105// loop over the outer event collection and the inner update collection,
106// i.e.:
107// foreach (var sse in _events) {
108// // get _updates from sse event
109// foreach (var update in _updates) { ... }
110// }
674e0f77Stephen Toub2 years ago111private IAsyncEnumerator<SseItem<byte[]>>? _events;
0ca4c062Jose Arriaga Maldonado1 years ago112private IEnumerator<U>? _updates;
113private readonly Func<SseItem<byte[]>, IEnumerable<U>> _deserializerFunc;
9f9f2936Jose Arriaga Maldonado2 years ago114
0ca4c062Jose Arriaga Maldonado1 years ago115private U? _current;
9f9f2936Jose Arriaga Maldonado2 years ago116private bool _started;
117
0ca4c062Jose Arriaga Maldonado1 years ago118public AsyncSseUpdateEnumerator(
119Func<SseItem<byte[]>, IEnumerable<U>> deserializerFunc,
120ClientResult page,
121CancellationToken cancellationToken)
9f9f2936Jose Arriaga Maldonado2 years ago122{
2ab1a942Jose Arriaga Maldonado1 years ago123Argument.AssertNotNull(page, nameof(page));
9f9f2936Jose Arriaga Maldonado2 years ago124
0ca4c062Jose Arriaga Maldonado1 years ago125_deserializerFunc = deserializerFunc;
2ab1a942Jose Arriaga Maldonado1 years ago126_response = page.GetRawResponse();
9f9f2936Jose Arriaga Maldonado2 years ago127_cancellationToken = cancellationToken;
128}
129
0ca4c062Jose Arriaga Maldonado1 years ago130U IAsyncEnumerator<U>.Current => _current!;
9f9f2936Jose Arriaga Maldonado2 years ago131
0ca4c062Jose Arriaga Maldonado1 years ago132async ValueTask<bool> IAsyncEnumerator<U>.MoveNextAsync()
9f9f2936Jose Arriaga Maldonado2 years ago133{
134if (_events is null && _started)
135{
0ca4c062Jose Arriaga Maldonado1 years ago136throw new ObjectDisposedException(nameof(AsyncSseUpdateEnumerator<U>));
9f9f2936Jose Arriaga Maldonado2 years ago137}
138
139_cancellationToken.ThrowIfCancellationRequested();
2ab1a942Jose Arriaga Maldonado1 years ago140_events ??= CreateEventEnumeratorAsync();
9f9f2936Jose Arriaga Maldonado2 years ago141_started = true;
142
143if (_updates is not null && _updates.MoveNext())
144{
145_current = _updates.Current;
146return true;
147}
148
149if (await _events.MoveNextAsync().ConfigureAwait(false))
150{
674e0f77Stephen Toub2 years ago151if (_events.Current.Data.AsSpan().SequenceEqual(TerminalData))
9f9f2936Jose Arriaga Maldonado2 years ago152{
153_current = default;
154return false;
155}
156
0ca4c062Jose Arriaga Maldonado1 years ago157_updates = _deserializerFunc
158.Invoke(_events.Current)
159.GetEnumerator();
9f9f2936Jose Arriaga Maldonado2 years ago160
161if (_updates.MoveNext())
162{
163_current = _updates.Current;
164return true;
165}
166}
167
168_current = default;
169return false;
170}
171
2ab1a942Jose Arriaga Maldonado1 years ago172private IAsyncEnumerator<SseItem<byte[]>> CreateEventEnumeratorAsync()
9f9f2936Jose Arriaga Maldonado2 years ago173{
2ab1a942Jose Arriaga Maldonado1 years ago174if (_response.ContentStream is null)
9f9f2936Jose Arriaga Maldonado2 years ago175{
176throw new InvalidOperationException("Unable to create result from response with null ContentStream");
177}
178
2ab1a942Jose Arriaga Maldonado1 years ago179IAsyncEnumerable<SseItem<byte[]>> enumerable = SseParser.Create(_response.ContentStream, (_, bytes) => bytes.ToArray()).EnumerateAsync();
9f9f2936Jose Arriaga Maldonado2 years ago180return enumerable.GetAsyncEnumerator(_cancellationToken);
181}
182
183public async ValueTask DisposeAsync()
184{
185await DisposeAsyncCore().ConfigureAwait(false);
186
187GC.SuppressFinalize(this);
188}
189
190private async ValueTask DisposeAsyncCore()
191{
192if (_events is not null)
193{
194await _events.DisposeAsync().ConfigureAwait(false);
195_events = null;
196
2ab1a942Jose Arriaga Maldonado1 years ago197// Dispose the response so we don't leave the network connection open.
198_response?.Dispose();
9f9f2936Jose Arriaga Maldonado2 years ago199}
200}
201}
202}