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.9.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

255lines · 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
fb7c8db4ShivangiReja10 months ago50public AsyncSseUpdateCollection(
51Func<Task<ClientResult>> sendRequestAsync,
52Func<JsonElement, BinaryData, ModelReaderWriterOptions, IEnumerable<T>> jsonMultiDeserializerFunc,
53CancellationToken cancellationToken)
54: this(
55sendRequestAsync,
56DeserializeSseToMultipleViaJson(jsonMultiDeserializerFunc),
57cancellationToken)
58{
59Argument.AssertNotNull(jsonMultiDeserializerFunc, nameof(jsonMultiDeserializerFunc));
60}
61
62public AsyncSseUpdateCollection(
63Func<Task<ClientResult>> sendRequestAsync,
64Func<JsonElement, BinaryData, ModelReaderWriterOptions, T> jsonSingleDeserializerFunc,
65CancellationToken cancellationToken)
66: this(
67sendRequestAsync,
68DeserializeSseToSingleViaJson(jsonSingleDeserializerFunc),
69cancellationToken)
70{
71Argument.AssertNotNull(jsonSingleDeserializerFunc, nameof(jsonSingleDeserializerFunc));
72}
73
0ca4c062Jose Arriaga Maldonado1 years ago74public AsyncSseUpdateCollection(
75Func<Task<ClientResult>> sendRequestAsync,
76Func<SseItem<byte[]>, IEnumerable<T>> eventDeserializerFunc,
2ab1a942Jose Arriaga Maldonado1 years ago77CancellationToken cancellationToken)
9f9f2936Jose Arriaga Maldonado2 years ago78{
2ab1a942Jose Arriaga Maldonado1 years ago79Argument.AssertNotNull(sendRequestAsync, nameof(sendRequestAsync));
0ca4c062Jose Arriaga Maldonado1 years ago80Argument.AssertNotNull(eventDeserializerFunc, nameof(eventDeserializerFunc));
9f9f2936Jose Arriaga Maldonado2 years ago81
2ab1a942Jose Arriaga Maldonado1 years ago82_sendRequestAsync = sendRequestAsync;
0ca4c062Jose Arriaga Maldonado1 years ago83_eventDeserializerFunc = eventDeserializerFunc;
2ab1a942Jose Arriaga Maldonado1 years ago84_cancellationToken = cancellationToken;
9f9f2936Jose Arriaga Maldonado2 years ago85}
86
2ab1a942Jose Arriaga Maldonado1 years ago87public override ContinuationToken? GetContinuationToken(ClientResult page)
88// Continuation is not supported for SSE streams.
89=> null;
90
91public 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.
95yield return await _sendRequestAsync();
96}
97
0ca4c062Jose Arriaga Maldonado1 years ago98protected async override IAsyncEnumerable<T> GetValuesFromPageAsync(ClientResult page)
9f9f2936Jose Arriaga Maldonado2 years ago99{
5dce104aJose Arriaga Maldonado1 years ago100await using IAsyncEnumerator<T> enumerator = new AsyncSseUpdateEnumerator<T>(_eventDeserializerFunc, page, _cancellationToken, AdditionalDisposalActions);
101
2ab1a942Jose Arriaga Maldonado1 years ago102while (await enumerator.MoveNextAsync().ConfigureAwait(false))
103{
104yield return enumerator.Current;
105}
9f9f2936Jose Arriaga Maldonado2 years ago106}
107
0ca4c062Jose Arriaga Maldonado1 years ago108[MethodImpl(MethodImplOptions.AggressiveInlining)]
109internal static Func<SseItem<byte[]>, IEnumerable<U>> DeserializeSseToMultipleViaJson<U>(
110Func<JsonElement, ModelReaderWriterOptions, IEnumerable<U>> jsonDeserializationFunc)
111{
112return (item) =>
113{
114using JsonDocument document = JsonDocument.Parse(item.Data);
115return jsonDeserializationFunc.Invoke(document.RootElement, ModelSerializationExtensions.WireOptions);
116};
117}
118
119[MethodImpl(MethodImplOptions.AggressiveInlining)]
120internal static Func<SseItem<byte[]>, IEnumerable<U>> DeserializeSseToSingleViaJson<U>(
121Func<JsonElement, ModelReaderWriterOptions, U> jsonSingleDeserializationFunc)
122=> DeserializeSseToMultipleViaJson<U>((e, o) => [jsonSingleDeserializationFunc.Invoke(e, o)]);
123
fb7c8db4ShivangiReja10 months ago124[MethodImpl(MethodImplOptions.AggressiveInlining)]
125internal static Func<SseItem<byte[]>, IEnumerable<U>> DeserializeSseToMultipleViaJson<U>(
126Func<JsonElement, BinaryData, ModelReaderWriterOptions, IEnumerable<U>> jsonDeserializationFunc)
127{
128return (item) =>
129{
130using JsonDocument document = JsonDocument.Parse(item.Data);
131return jsonDeserializationFunc.Invoke(document.RootElement, BinaryData.FromBytes(item.Data), ModelSerializationExtensions.WireOptions);
132};
133}
134
135[MethodImpl(MethodImplOptions.AggressiveInlining)]
136internal static Func<SseItem<byte[]>, IEnumerable<U>> DeserializeSseToSingleViaJson<U>(
137Func<JsonElement, BinaryData, ModelReaderWriterOptions, U> jsonSingleDeserializationFunc)
138=> DeserializeSseToMultipleViaJson<U>((e, d, o) => [jsonSingleDeserializationFunc.Invoke(e, d, o)]);
139
0ca4c062Jose Arriaga Maldonado1 years ago140private sealed class AsyncSseUpdateEnumerator<U> : IAsyncEnumerator<U>
9f9f2936Jose Arriaga Maldonado2 years ago141{
674e0f77Stephen Toub2 years ago142private static ReadOnlySpan<byte> TerminalData => "[DONE]"u8;
9f9f2936Jose Arriaga Maldonado2 years ago143
5dce104aJose Arriaga Maldonado1 years ago144private List<Action> _additionalDisposalActions;
145
9f9f2936Jose Arriaga Maldonado2 years ago146private readonly CancellationToken _cancellationToken;
2ab1a942Jose Arriaga Maldonado1 years ago147private readonly PipelineResponse _response;
9f9f2936Jose Arriaga Maldonado2 years ago148
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// }
674e0f77Stephen Toub2 years ago156private IAsyncEnumerator<SseItem<byte[]>>? _events;
0ca4c062Jose Arriaga Maldonado1 years ago157private IEnumerator<U>? _updates;
158private readonly Func<SseItem<byte[]>, IEnumerable<U>> _deserializerFunc;
9f9f2936Jose Arriaga Maldonado2 years ago159
0ca4c062Jose Arriaga Maldonado1 years ago160private U? _current;
9f9f2936Jose Arriaga Maldonado2 years ago161private bool _started;
162
0ca4c062Jose Arriaga Maldonado1 years ago163public AsyncSseUpdateEnumerator(
164Func<SseItem<byte[]>, IEnumerable<U>> deserializerFunc,
165ClientResult page,
5dce104aJose Arriaga Maldonado1 years ago166CancellationToken cancellationToken,
167List<Action> additionalDisposalActions)
9f9f2936Jose Arriaga Maldonado2 years ago168{
2ab1a942Jose Arriaga Maldonado1 years ago169Argument.AssertNotNull(page, nameof(page));
9f9f2936Jose Arriaga Maldonado2 years ago170
0ca4c062Jose Arriaga Maldonado1 years ago171_deserializerFunc = deserializerFunc;
2ab1a942Jose Arriaga Maldonado1 years ago172_response = page.GetRawResponse();
9f9f2936Jose Arriaga Maldonado2 years ago173_cancellationToken = cancellationToken;
5dce104aJose Arriaga Maldonado1 years ago174_additionalDisposalActions = additionalDisposalActions;
9f9f2936Jose Arriaga Maldonado2 years ago175}
176
0ca4c062Jose Arriaga Maldonado1 years ago177U IAsyncEnumerator<U>.Current => _current!;
9f9f2936Jose Arriaga Maldonado2 years ago178
0ca4c062Jose Arriaga Maldonado1 years ago179async ValueTask<bool> IAsyncEnumerator<U>.MoveNextAsync()
9f9f2936Jose Arriaga Maldonado2 years ago180{
181if (_events is null && _started)
182{
0ca4c062Jose Arriaga Maldonado1 years ago183throw new ObjectDisposedException(nameof(AsyncSseUpdateEnumerator<U>));
9f9f2936Jose Arriaga Maldonado2 years ago184}
185
186_cancellationToken.ThrowIfCancellationRequested();
2ab1a942Jose Arriaga Maldonado1 years ago187_events ??= CreateEventEnumeratorAsync();
9f9f2936Jose Arriaga Maldonado2 years ago188_started = true;
189
190if (_updates is not null && _updates.MoveNext())
191{
192_current = _updates.Current;
193return true;
194}
195
196if (await _events.MoveNextAsync().ConfigureAwait(false))
197{
674e0f77Stephen Toub2 years ago198if (_events.Current.Data.AsSpan().SequenceEqual(TerminalData))
9f9f2936Jose Arriaga Maldonado2 years ago199{
200_current = default;
201return false;
202}
203
0ca4c062Jose Arriaga Maldonado1 years ago204_updates = _deserializerFunc
205.Invoke(_events.Current)
206.GetEnumerator();
9f9f2936Jose Arriaga Maldonado2 years ago207
208if (_updates.MoveNext())
209{
210_current = _updates.Current;
211return true;
212}
213}
214
215_current = default;
216return false;
217}
218
2ab1a942Jose Arriaga Maldonado1 years ago219private IAsyncEnumerator<SseItem<byte[]>> CreateEventEnumeratorAsync()
9f9f2936Jose Arriaga Maldonado2 years ago220{
2ab1a942Jose Arriaga Maldonado1 years ago221if (_response.ContentStream is null)
9f9f2936Jose Arriaga Maldonado2 years ago222{
223throw new InvalidOperationException("Unable to create result from response with null ContentStream");
224}
225
2ab1a942Jose Arriaga Maldonado1 years ago226IAsyncEnumerable<SseItem<byte[]>> enumerable = SseParser.Create(_response.ContentStream, (_, bytes) => bytes.ToArray()).EnumerateAsync();
9f9f2936Jose Arriaga Maldonado2 years ago227return enumerable.GetAsyncEnumerator(_cancellationToken);
228}
229
230public async ValueTask DisposeAsync()
231{
232await DisposeAsyncCore().ConfigureAwait(false);
233
234GC.SuppressFinalize(this);
235}
236
237private async ValueTask DisposeAsyncCore()
238{
239if (_events is not null)
240{
241await _events.DisposeAsync().ConfigureAwait(false);
242_events = null;
243
2ab1a942Jose Arriaga Maldonado1 years ago244// Dispose the response so we don't leave the network connection open.
245_response?.Dispose();
9f9f2936Jose Arriaga Maldonado2 years ago246}
5dce104aJose Arriaga Maldonado1 years ago247
248foreach (Action additionalDisposalAction in _additionalDisposalActions ?? [])
249{
250additionalDisposalAction.Invoke();
251}
252_additionalDisposalActions?.Clear();
9f9f2936Jose Arriaga Maldonado2 years ago253}
254}
255}