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-beta.2

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 · modecode

RawDownload
Latest commit unavailable.
1using System;
2using System.ClientModel;
3using System.ClientModel.Primitives;
4using System.Collections.Generic;
5using System.Linq;
6using System.Net.ServerSentEvents;
7using 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>
17internal class AsyncStreamingUpdateCollection : AsyncCollectionResult<StreamingUpdate>
18{
19 private readonly Func<Task<ClientResult>> _sendRequestAsync;
20 private readonly CancellationToken _cancellationToken;
21
22 public AsyncStreamingUpdateCollection(Func<Task<ClientResult>> sendRequestAsync,
23 CancellationToken cancellationToken)
24 {
25 Argument.AssertNotNull(sendRequestAsync, nameof(sendRequestAsync));
26
27 _sendRequestAsync = sendRequestAsync;
28 _cancellationToken = cancellationToken;
29 }
30
31 public override ContinuationToken? GetContinuationToken(ClientResult page)
32 // Continuation is not supported for SSE streams.
33 => null;
34
35 public async override IAsyncEnumerable<ClientResult> GetRawPagesAsync()
36 {
37 // We don't currently support resuming a dropped connection from the
38 // last received event, so the response collection has a single element.
39 yield return await _sendRequestAsync();
40 }
41
42 protected async override IAsyncEnumerable<StreamingUpdate> GetValuesFromPageAsync(ClientResult page)
43 {
44 await using IAsyncEnumerator<StreamingUpdate> enumerator = new AsyncStreamingUpdateEnumerator(page, _cancellationToken);
45 while (await enumerator.MoveNextAsync().ConfigureAwait(false))
46 {
47 yield return enumerator.Current;
48 }
49 }
50
51 private sealed class AsyncStreamingUpdateEnumerator : IAsyncEnumerator<StreamingUpdate>
52 {
53 private static ReadOnlySpan<byte> TerminalData => "[DONE]"u8;
54
55 private readonly CancellationToken _cancellationToken;
56 private readonly PipelineResponse _response;
57
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 // }
65 private IAsyncEnumerator<SseItem<byte[]>>? _events;
66 private IEnumerator<StreamingUpdate>? _updates;
67
68 private StreamingUpdate? _current;
69 private bool _started;
70
71 public AsyncStreamingUpdateEnumerator(ClientResult page, CancellationToken cancellationToken)
72 {
73 Argument.AssertNotNull(page, nameof(page));
74
75 _response = page.GetRawResponse();
76 _cancellationToken = cancellationToken;
77 }
78
79 StreamingUpdate IAsyncEnumerator<StreamingUpdate>.Current
80 => _current!;
81
82 async ValueTask<bool> IAsyncEnumerator<StreamingUpdate>.MoveNextAsync()
83 {
84 if (_events is null && _started)
85 {
86 throw new ObjectDisposedException(nameof(AsyncStreamingUpdateEnumerator));
87 }
88
89 _cancellationToken.ThrowIfCancellationRequested();
90 _events ??= CreateEventEnumeratorAsync();
91 _started = true;
92
93 if (_updates is not null && _updates.MoveNext())
94 {
95 _current = _updates.Current;
96 return true;
97 }
98
99 if (await _events.MoveNextAsync().ConfigureAwait(false))
100 {
101 if (_events.Current.Data.AsSpan().SequenceEqual(TerminalData))
102 {
103 _current = default;
104 return false;
105 }
106
107 var updates = StreamingUpdate.FromEvent(_events.Current);
108 _updates = updates.GetEnumerator();
109
110 if (_updates.MoveNext())
111 {
112 _current = _updates.Current;
113 return true;
114 }
115 }
116
117 _current = default;
118 return false;
119 }
120
121 private IAsyncEnumerator<SseItem<byte[]>> CreateEventEnumeratorAsync()
122 {
123 if (_response.ContentStream is null)
124 {
125 throw new InvalidOperationException("Unable to create result from response with null ContentStream");
126 }
127
128 IAsyncEnumerable<SseItem<byte[]>> enumerable = SseParser.Create(_response.ContentStream, (_, bytes) => bytes.ToArray()).EnumerateAsync();
129 return enumerable.GetAsyncEnumerator(_cancellationToken);
130 }
131
132 public async ValueTask DisposeAsync()
133 {
134 await DisposeAsyncCore().ConfigureAwait(false);
135
136 GC.SuppressFinalize(this);
137 }
138
139 private async ValueTask DisposeAsyncCore()
140 {
141 if (_events is not null)
142 {
143 await _events.DisposeAsync().ConfigureAwait(false);
144 _events = null;
145
146 // Dispose the response so we don't leave the network connection open.
147 _response?.Dispose();
148 }
149 }
150 }
151}
152