openai/openai-dotnet

Public

mirrored from https://github.com/openai/openai-dotnetAvailable

CodeCommitsIssuesPull requestsActionsInsightsSecurity
6b3c8b0a3d3208218dc3cbcd295502024988a956

Branches

Tags

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

Clone

HTTPS

Download ZIP

src/Utility/SseUpdateCollection.cs

228lines · modecode

1using System;
2using System.ClientModel;
3using System.ClientModel.Primitives;
4using System.Collections;
5using System.Collections.Generic;
6using System.Net.ServerSentEvents;
7using System.Text.Json;
8using System.Threading;
9
10#nullable enable
11
12namespace OpenAI;
13
14/// <summary>
15/// Implementation of collection abstraction over streaming updates.
16/// </summary>
17internal class SseUpdateCollection<T> : CollectionResult<T>
18{
19 private readonly Func<ClientResult> _sendRequestFunc;
20 private readonly Func<SseItem<byte[]>, IEnumerable<T>> _eventDeserializerFunc;
21 private readonly CancellationToken _cancellationToken;
22
23 public List<Action> AdditionalDisposalActions { get; } = [];
24
25 public SseUpdateCollection(
26 Func<ClientResult> sendRequestFunc,
27 Func<JsonElement, ModelReaderWriterOptions, IEnumerable<T>> jsonMultiDeserializerFunc,
28 CancellationToken cancellationToken)
29 : this(
30 sendRequestFunc,
31 AsyncSseUpdateCollection<T>.DeserializeSseToMultipleViaJson(jsonMultiDeserializerFunc),
32 cancellationToken)
33
34 {
35 Argument.AssertNotNull(jsonMultiDeserializerFunc, nameof(jsonMultiDeserializerFunc));
36 }
37
38 public SseUpdateCollection(
39 Func<ClientResult> sendRequestFunc,
40 Func<JsonElement, ModelReaderWriterOptions, T> jsonSingleDeserializerFunc,
41 CancellationToken cancellationToken)
42 : this(
43 sendRequestFunc,
44 AsyncSseUpdateCollection<T>.DeserializeSseToSingleViaJson(jsonSingleDeserializerFunc),
45 cancellationToken)
46 {
47 Argument.AssertNotNull(jsonSingleDeserializerFunc, nameof(jsonSingleDeserializerFunc));
48 }
49
50 public SseUpdateCollection(
51 Func<ClientResult> sendRequestFunc,
52 Func<JsonElement, BinaryData, ModelReaderWriterOptions, IEnumerable<T>> jsonMultiDeserializerFunc,
53 CancellationToken cancellationToken)
54 : this(
55 sendRequestFunc,
56 AsyncSseUpdateCollection<T>.DeserializeSseToMultipleViaJson(jsonMultiDeserializerFunc),
57 cancellationToken)
58
59 {
60 Argument.AssertNotNull(jsonMultiDeserializerFunc, nameof(jsonMultiDeserializerFunc));
61 }
62
63 public SseUpdateCollection(
64 Func<ClientResult> sendRequestFunc,
65 Func<JsonElement, BinaryData, ModelReaderWriterOptions, T> jsonSingleDeserializerFunc,
66 CancellationToken cancellationToken)
67 : this(
68 sendRequestFunc,
69 AsyncSseUpdateCollection<T>.DeserializeSseToSingleViaJson(jsonSingleDeserializerFunc),
70 cancellationToken)
71 {
72 Argument.AssertNotNull(jsonSingleDeserializerFunc, nameof(jsonSingleDeserializerFunc));
73 }
74
75 public SseUpdateCollection(
76 Func<ClientResult> sendRequestFunc,
77 Func<SseItem<byte[]>, IEnumerable<T>> eventDeserializerFunc,
78 CancellationToken cancellationToken)
79 {
80 Argument.AssertNotNull(sendRequestFunc, nameof(sendRequestFunc));
81 Argument.AssertNotNull(eventDeserializerFunc, nameof(eventDeserializerFunc));
82
83 _sendRequestFunc = sendRequestFunc;
84 _eventDeserializerFunc = eventDeserializerFunc;
85 _cancellationToken = cancellationToken;
86 }
87
88 public override ContinuationToken? GetContinuationToken(ClientResult page)
89 // Continuation is not supported for SSE streams.
90 => null;
91
92 public override IEnumerable<ClientResult> GetRawPages()
93 {
94 // We don't currently support resuming a dropped connection from the
95 // last received event, so the response collection has a single element.
96 yield return _sendRequestFunc();
97 }
98
99 protected override IEnumerable<T> GetValuesFromPage(ClientResult page)
100 {
101 using IEnumerator<T> enumerator = new SseUpdateEnumerator<T>(_eventDeserializerFunc, page, _cancellationToken, AdditionalDisposalActions);
102 while (enumerator.MoveNext())
103 {
104 yield return enumerator.Current;
105 }
106 }
107
108 private sealed class SseUpdateEnumerator<U> : IEnumerator<U>
109 {
110 private static ReadOnlySpan<byte> TerminalData => "[DONE]"u8;
111
112 private List<Action> _additionalDisposalActions;
113
114 private readonly CancellationToken _cancellationToken;
115 private readonly PipelineResponse _response;
116
117 // These enumerators represent what is effectively a doubly-nested
118 // loop over the outer event collection and the inner update collection,
119 // i.e.:
120 // foreach (var sse in _events) {
121 // // get _updates from sse event
122 // foreach (var update in _updates) { ... }
123 // }
124 private IEnumerator<SseItem<byte[]>>? _events;
125 private IEnumerator<U>? _updates;
126 private readonly Func<SseItem<byte[]>, IEnumerable<U>> _eventDeserializerFunc;
127
128 private U? _current;
129 private bool _started;
130
131 public SseUpdateEnumerator(
132 Func<SseItem<byte[]>, IEnumerable<U>> eventDeserializerFunc,
133 ClientResult page,
134 CancellationToken cancellationToken,
135 List<Action> additionalDisposalActions)
136 {
137 Argument.AssertNotNull(eventDeserializerFunc, nameof(eventDeserializerFunc));
138 Argument.AssertNotNull(page, nameof(page));
139
140 _eventDeserializerFunc = eventDeserializerFunc;
141 _response = page.GetRawResponse();
142 _cancellationToken = cancellationToken;
143 _additionalDisposalActions = additionalDisposalActions;
144 }
145
146 U IEnumerator<U>.Current => _current!;
147
148 object IEnumerator.Current => _current!;
149
150 public bool MoveNext()
151 {
152 if (_events is null && _started)
153 {
154 throw new ObjectDisposedException(typeof(U).Name);
155 }
156
157 _cancellationToken.ThrowIfCancellationRequested();
158 _events ??= CreateEventEnumerator();
159 _started = true;
160
161 if (_updates is not null && _updates.MoveNext())
162 {
163 _current = _updates.Current;
164 return true;
165 }
166
167 if (_events.MoveNext())
168 {
169 if (_events.Current.Data.AsSpan().SequenceEqual(TerminalData))
170 {
171 _current = default;
172 return false;
173 }
174
175 _updates = _eventDeserializerFunc.Invoke(_events.Current).GetEnumerator();
176
177 if (_updates.MoveNext())
178 {
179 _current = _updates.Current;
180 return true;
181 }
182 }
183
184 _current = default;
185 return false;
186 }
187
188 private IEnumerator<SseItem<byte[]>> CreateEventEnumerator()
189 {
190 if (_response.ContentStream is null)
191 {
192 throw new InvalidOperationException("Unable to create result from response with null ContentStream");
193 }
194
195 IEnumerable<SseItem<byte[]>> enumerable = SseParser.Create(_response.ContentStream, (_, bytes) => bytes.ToArray()).Enumerate();
196 return enumerable.GetEnumerator();
197 }
198
199 public void Reset()
200 {
201 throw new NotSupportedException("Cannot seek back in an SSE stream.");
202 }
203
204 public void Dispose()
205 {
206 Dispose(true);
207 GC.SuppressFinalize(this);
208 }
209
210 private void Dispose(bool disposing)
211 {
212 if (disposing && _events is not null)
213 {
214 _events.Dispose();
215 _events = null;
216
217 // Dispose the response so we don't leave the network connection open.
218 _response?.Dispose();
219 }
220
221 foreach (Action additionalDisposalAction in _additionalDisposalActions ?? [])
222 {
223 additionalDisposalAction?.Invoke();
224 }
225 _additionalDisposalActions?.Clear();
226 }
227 }
228}
229