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.0.0-beta.12

Branches

Tags

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

Clone

HTTPS

Download ZIP

openai-dotnet/src/Custom/Assistants/Streaming

src/Custom/Assistants/Streaming/AsyncStreamingUpdateCollection.cs

151lines · 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.Linq;
6using System.Net.ServerSentEvents;
9f9f2936Jose Arriaga Maldonado2 years ago7using System.Threading;
8using System.Threading.Tasks;
9
10#nullable enable
11
12namespace OpenAI.Assistants;
13
14/// <summary>
15/// Implementation of collection abstraction over streaming assistant updates.
16/// </summary>
7bdecfd8Anne Thompson2 years ago17internal class AsyncStreamingUpdateCollection : AsyncCollectionResult<StreamingUpdate>
9f9f2936Jose Arriaga Maldonado2 years ago18{
2ab1a942Jose Arriaga Maldonado1 years ago19private readonly Func<Task<ClientResult>> _sendRequestAsync;
20private readonly CancellationToken _cancellationToken;
9f9f2936Jose Arriaga Maldonado2 years ago21
2ab1a942Jose Arriaga Maldonado1 years ago22public AsyncStreamingUpdateCollection(Func<Task<ClientResult>> sendRequestAsync,
23CancellationToken cancellationToken)
9f9f2936Jose Arriaga Maldonado2 years ago24{
2ab1a942Jose Arriaga Maldonado1 years ago25Argument.AssertNotNull(sendRequestAsync, nameof(sendRequestAsync));
9f9f2936Jose Arriaga Maldonado2 years ago26
2ab1a942Jose Arriaga Maldonado1 years ago27_sendRequestAsync = sendRequestAsync;
28_cancellationToken = cancellationToken;
9f9f2936Jose Arriaga Maldonado2 years ago29}
30
2ab1a942Jose Arriaga Maldonado1 years ago31public override ContinuationToken? GetContinuationToken(ClientResult page)
32// Continuation is not supported for SSE streams.
33=> null;
34
35public async override IAsyncEnumerable<ClientResult> GetRawPagesAsync()
9f9f2936Jose Arriaga Maldonado2 years ago36{
2ab1a942Jose Arriaga Maldonado1 years ago37// We don't currently support resuming a dropped connection from the
38// last received event, so the response collection has a single element.
39yield return await _sendRequestAsync();
40}
41
42protected async override IAsyncEnumerable<StreamingUpdate> GetValuesFromPageAsync(ClientResult page)
43{
44await using IAsyncEnumerator<StreamingUpdate> enumerator = new AsyncStreamingUpdateEnumerator(page, _cancellationToken);
45while (await enumerator.MoveNextAsync().ConfigureAwait(false))
46{
47yield return enumerator.Current;
48}
9f9f2936Jose Arriaga Maldonado2 years ago49}
50
51private sealed class AsyncStreamingUpdateEnumerator : IAsyncEnumerator<StreamingUpdate>
52{
674e0f77Stephen Toub2 years ago53private static ReadOnlySpan<byte> TerminalData => "[DONE]"u8;
9f9f2936Jose Arriaga Maldonado2 years ago54
55private readonly CancellationToken _cancellationToken;
2ab1a942Jose Arriaga Maldonado1 years ago56private readonly PipelineResponse _response;
9f9f2936Jose Arriaga Maldonado2 years ago57
58// These enumerators represent what is effectively a doubly-nested
59// loop over the outer event collection and the inner update collection,
60// i.e.:
61// foreach (var sse in _events) {
62// // get _updates from sse event
63// foreach (var update in _updates) { ... }
64// }
674e0f77Stephen Toub2 years ago65private IAsyncEnumerator<SseItem<byte[]>>? _events;
9f9f2936Jose Arriaga Maldonado2 years ago66private IEnumerator<StreamingUpdate>? _updates;
67
68private StreamingUpdate? _current;
69private bool _started;
70
2ab1a942Jose Arriaga Maldonado1 years ago71public AsyncStreamingUpdateEnumerator(ClientResult page, CancellationToken cancellationToken)
9f9f2936Jose Arriaga Maldonado2 years ago72{
2ab1a942Jose Arriaga Maldonado1 years ago73Argument.AssertNotNull(page, nameof(page));
9f9f2936Jose Arriaga Maldonado2 years ago74
2ab1a942Jose Arriaga Maldonado1 years ago75_response = page.GetRawResponse();
9f9f2936Jose Arriaga Maldonado2 years ago76_cancellationToken = cancellationToken;
77}
78
79StreamingUpdate IAsyncEnumerator<StreamingUpdate>.Current
80=> _current!;
81
82async ValueTask<bool> IAsyncEnumerator<StreamingUpdate>.MoveNextAsync()
83{
84if (_events is null && _started)
85{
86throw new ObjectDisposedException(nameof(AsyncStreamingUpdateEnumerator));
87}
88
89_cancellationToken.ThrowIfCancellationRequested();
2ab1a942Jose Arriaga Maldonado1 years ago90_events ??= CreateEventEnumeratorAsync();
9f9f2936Jose Arriaga Maldonado2 years ago91_started = true;
92
93if (_updates is not null && _updates.MoveNext())
94{
95_current = _updates.Current;
96return true;
97}
98
99if (await _events.MoveNextAsync().ConfigureAwait(false))
100{
674e0f77Stephen Toub2 years ago101if (_events.Current.Data.AsSpan().SequenceEqual(TerminalData))
9f9f2936Jose Arriaga Maldonado2 years ago102{
103_current = default;
104return false;
105}
106
107var updates = StreamingUpdate.FromEvent(_events.Current);
108_updates = updates.GetEnumerator();
109
110if (_updates.MoveNext())
111{
112_current = _updates.Current;
113return true;
114}
115}
116
117_current = default;
118return false;
119}
120
2ab1a942Jose Arriaga Maldonado1 years ago121private IAsyncEnumerator<SseItem<byte[]>> CreateEventEnumeratorAsync()
9f9f2936Jose Arriaga Maldonado2 years ago122{
2ab1a942Jose Arriaga Maldonado1 years ago123if (_response.ContentStream is null)
9f9f2936Jose Arriaga Maldonado2 years ago124{
125throw new InvalidOperationException("Unable to create result from response with null ContentStream");
126}
127
2ab1a942Jose Arriaga Maldonado1 years ago128IAsyncEnumerable<SseItem<byte[]>> enumerable = SseParser.Create(_response.ContentStream, (_, bytes) => bytes.ToArray()).EnumerateAsync();
9f9f2936Jose Arriaga Maldonado2 years ago129return enumerable.GetAsyncEnumerator(_cancellationToken);
130}
131
132public async ValueTask DisposeAsync()
133{
134await DisposeAsyncCore().ConfigureAwait(false);
135
136GC.SuppressFinalize(this);
137}
138
139private async ValueTask DisposeAsyncCore()
140{
141if (_events is not null)
142{
143await _events.DisposeAsync().ConfigureAwait(false);
144_events = null;
145
2ab1a942Jose Arriaga Maldonado1 years ago146// Dispose the response so we don't leave the network connection open.
147_response?.Dispose();
9f9f2936Jose Arriaga Maldonado2 years ago148}
149}
150}
151}