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
achandmsft-patch-2

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 · 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
5dce104aJose Arriaga Maldonado1 years ago24public List<Action> AdditionalDisposalActions { get; } = [];
25
0ca4c062Jose Arriaga Maldonado1 years ago26public AsyncSseUpdateCollection(
2ab1a942Jose Arriaga Maldonado1 years ago27Func<Task<ClientResult>> sendRequestAsync,
0ca4c062Jose Arriaga Maldonado1 years ago28Func<JsonElement, ModelReaderWriterOptions, IEnumerable<T>> jsonMultiDeserializerFunc,
29CancellationToken cancellationToken)
30: this(
31sendRequestAsync,
32DeserializeSseToMultipleViaJson(jsonMultiDeserializerFunc),
33cancellationToken)
34{
35Argument.AssertNotNull(jsonMultiDeserializerFunc, nameof(jsonMultiDeserializerFunc));
36}
37
38public AsyncSseUpdateCollection(
39Func<Task<ClientResult>> sendRequestAsync,
40Func<JsonElement, ModelReaderWriterOptions, T> jsonSingleDeserializerFunc,
41CancellationToken cancellationToken)
42: this(
43sendRequestAsync,
44DeserializeSseToSingleViaJson(jsonSingleDeserializerFunc),
45cancellationToken)
46{
47Argument.AssertNotNull(jsonSingleDeserializerFunc, nameof(jsonSingleDeserializerFunc));
48}
49
50public AsyncSseUpdateCollection(
51Func<Task<ClientResult>> sendRequestAsync,
52Func<SseItem<byte[]>, IEnumerable<T>> eventDeserializerFunc,
2ab1a942Jose Arriaga Maldonado1 years ago53CancellationToken cancellationToken)
9f9f2936Jose Arriaga Maldonado2 years ago54{
2ab1a942Jose Arriaga Maldonado1 years ago55Argument.AssertNotNull(sendRequestAsync, nameof(sendRequestAsync));
0ca4c062Jose Arriaga Maldonado1 years ago56Argument.AssertNotNull(eventDeserializerFunc, nameof(eventDeserializerFunc));
9f9f2936Jose Arriaga Maldonado2 years ago57
2ab1a942Jose Arriaga Maldonado1 years ago58_sendRequestAsync = sendRequestAsync;
0ca4c062Jose Arriaga Maldonado1 years ago59_eventDeserializerFunc = eventDeserializerFunc;
2ab1a942Jose Arriaga Maldonado1 years ago60_cancellationToken = cancellationToken;
9f9f2936Jose Arriaga Maldonado2 years ago61}
62
2ab1a942Jose Arriaga Maldonado1 years ago63public override ContinuationToken? GetContinuationToken(ClientResult page)
64// Continuation is not supported for SSE streams.
65=> null;
66
67public 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.
71yield return await _sendRequestAsync();
72}
73
0ca4c062Jose Arriaga Maldonado1 years ago74protected async override IAsyncEnumerable<T> GetValuesFromPageAsync(ClientResult page)
9f9f2936Jose Arriaga Maldonado2 years ago75{
5dce104aJose Arriaga Maldonado1 years ago76await using IAsyncEnumerator<T> enumerator = new AsyncSseUpdateEnumerator<T>(_eventDeserializerFunc, page, _cancellationToken, AdditionalDisposalActions);
77
2ab1a942Jose Arriaga Maldonado1 years ago78while (await enumerator.MoveNextAsync().ConfigureAwait(false))
79{
80yield return enumerator.Current;
81}
9f9f2936Jose Arriaga Maldonado2 years ago82}
83
0ca4c062Jose Arriaga Maldonado1 years ago84[MethodImpl(MethodImplOptions.AggressiveInlining)]
85internal static Func<SseItem<byte[]>, IEnumerable<U>> DeserializeSseToMultipleViaJson<U>(
86Func<JsonElement, ModelReaderWriterOptions, IEnumerable<U>> jsonDeserializationFunc)
87{
88return (item) =>
89{
90using JsonDocument document = JsonDocument.Parse(item.Data);
91return jsonDeserializationFunc.Invoke(document.RootElement, ModelSerializationExtensions.WireOptions);
92};
93}
94
95[MethodImpl(MethodImplOptions.AggressiveInlining)]
96internal static Func<SseItem<byte[]>, IEnumerable<U>> DeserializeSseToSingleViaJson<U>(
97Func<JsonElement, ModelReaderWriterOptions, U> jsonSingleDeserializationFunc)
98=> DeserializeSseToMultipleViaJson<U>((e, o) => [jsonSingleDeserializationFunc.Invoke(e, o)]);
99
100private sealed class AsyncSseUpdateEnumerator<U> : IAsyncEnumerator<U>
9f9f2936Jose Arriaga Maldonado2 years ago101{
674e0f77Stephen Toub2 years ago102private static ReadOnlySpan<byte> TerminalData => "[DONE]"u8;
9f9f2936Jose Arriaga Maldonado2 years ago103
5dce104aJose Arriaga Maldonado1 years ago104private List<Action> _additionalDisposalActions;
105
9f9f2936Jose Arriaga Maldonado2 years ago106private readonly CancellationToken _cancellationToken;
2ab1a942Jose Arriaga Maldonado1 years ago107private readonly PipelineResponse _response;
9f9f2936Jose Arriaga Maldonado2 years ago108
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// }
674e0f77Stephen Toub2 years ago116private IAsyncEnumerator<SseItem<byte[]>>? _events;
0ca4c062Jose Arriaga Maldonado1 years ago117private IEnumerator<U>? _updates;
118private readonly Func<SseItem<byte[]>, IEnumerable<U>> _deserializerFunc;
9f9f2936Jose Arriaga Maldonado2 years ago119
0ca4c062Jose Arriaga Maldonado1 years ago120private U? _current;
9f9f2936Jose Arriaga Maldonado2 years ago121private bool _started;
122
0ca4c062Jose Arriaga Maldonado1 years ago123public AsyncSseUpdateEnumerator(
124Func<SseItem<byte[]>, IEnumerable<U>> deserializerFunc,
125ClientResult page,
5dce104aJose Arriaga Maldonado1 years ago126CancellationToken cancellationToken,
127List<Action> additionalDisposalActions)
9f9f2936Jose Arriaga Maldonado2 years ago128{
2ab1a942Jose Arriaga Maldonado1 years ago129Argument.AssertNotNull(page, nameof(page));
9f9f2936Jose Arriaga Maldonado2 years ago130
0ca4c062Jose Arriaga Maldonado1 years ago131_deserializerFunc = deserializerFunc;
2ab1a942Jose Arriaga Maldonado1 years ago132_response = page.GetRawResponse();
9f9f2936Jose Arriaga Maldonado2 years ago133_cancellationToken = cancellationToken;
5dce104aJose Arriaga Maldonado1 years ago134_additionalDisposalActions = additionalDisposalActions;
9f9f2936Jose Arriaga Maldonado2 years ago135}
136
0ca4c062Jose Arriaga Maldonado1 years ago137U IAsyncEnumerator<U>.Current => _current!;
9f9f2936Jose Arriaga Maldonado2 years ago138
0ca4c062Jose Arriaga Maldonado1 years ago139async ValueTask<bool> IAsyncEnumerator<U>.MoveNextAsync()
9f9f2936Jose Arriaga Maldonado2 years ago140{
141if (_events is null && _started)
142{
0ca4c062Jose Arriaga Maldonado1 years ago143throw new ObjectDisposedException(nameof(AsyncSseUpdateEnumerator<U>));
9f9f2936Jose Arriaga Maldonado2 years ago144}
145
146_cancellationToken.ThrowIfCancellationRequested();
2ab1a942Jose Arriaga Maldonado1 years ago147_events ??= CreateEventEnumeratorAsync();
9f9f2936Jose Arriaga Maldonado2 years ago148_started = true;
149
150if (_updates is not null && _updates.MoveNext())
151{
152_current = _updates.Current;
153return true;
154}
155
156if (await _events.MoveNextAsync().ConfigureAwait(false))
157{
674e0f77Stephen Toub2 years ago158if (_events.Current.Data.AsSpan().SequenceEqual(TerminalData))
9f9f2936Jose Arriaga Maldonado2 years ago159{
160_current = default;
161return false;
162}
163
0ca4c062Jose Arriaga Maldonado1 years ago164_updates = _deserializerFunc
165.Invoke(_events.Current)
166.GetEnumerator();
9f9f2936Jose Arriaga Maldonado2 years ago167
168if (_updates.MoveNext())
169{
170_current = _updates.Current;
171return true;
172}
173}
174
175_current = default;
176return false;
177}
178
2ab1a942Jose Arriaga Maldonado1 years ago179private IAsyncEnumerator<SseItem<byte[]>> CreateEventEnumeratorAsync()
9f9f2936Jose Arriaga Maldonado2 years ago180{
2ab1a942Jose Arriaga Maldonado1 years ago181if (_response.ContentStream is null)
9f9f2936Jose Arriaga Maldonado2 years ago182{
183throw new InvalidOperationException("Unable to create result from response with null ContentStream");
184}
185
2ab1a942Jose Arriaga Maldonado1 years ago186IAsyncEnumerable<SseItem<byte[]>> enumerable = SseParser.Create(_response.ContentStream, (_, bytes) => bytes.ToArray()).EnumerateAsync();
9f9f2936Jose Arriaga Maldonado2 years ago187return enumerable.GetAsyncEnumerator(_cancellationToken);
188}
189
190public async ValueTask DisposeAsync()
191{
192await DisposeAsyncCore().ConfigureAwait(false);
193
194GC.SuppressFinalize(this);
195}
196
197private async ValueTask DisposeAsyncCore()
198{
199if (_events is not null)
200{
201await _events.DisposeAsync().ConfigureAwait(false);
202_events = null;
203
2ab1a942Jose Arriaga Maldonado1 years ago204// Dispose the response so we don't leave the network connection open.
205_response?.Dispose();
9f9f2936Jose Arriaga Maldonado2 years ago206}
5dce104aJose Arriaga Maldonado1 years ago207
208foreach (Action additionalDisposalAction in _additionalDisposalActions ?? [])
209{
210additionalDisposalAction.Invoke();
211}
212_additionalDisposalActions?.Clear();
9f9f2936Jose Arriaga Maldonado2 years ago213}
214}
215}