openai/openai-dotnet

Public

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

CodeCommitsIssuesPull requestsActionsInsightsSecurity
OpenAI_2.2.0

Branches

Tags

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

Clone

HTTPS

Download ZIP

src/Utility/SseUpdateCollection.cs

203lines · 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<SseItem<byte[]>, IEnumerable<T>> eventDeserializerFunc,
53 CancellationToken cancellationToken)
54 {
55 Argument.AssertNotNull(sendRequestFunc, nameof(sendRequestFunc));
56 Argument.AssertNotNull(eventDeserializerFunc, nameof(eventDeserializerFunc));
57
58 _sendRequestFunc = sendRequestFunc;
59 _eventDeserializerFunc = eventDeserializerFunc;
60 _cancellationToken = cancellationToken;
61 }
62
63 public override ContinuationToken? GetContinuationToken(ClientResult page)
64 // Continuation is not supported for SSE streams.
65 => null;
66
67 public override IEnumerable<ClientResult> GetRawPages()
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.
71 yield return _sendRequestFunc();
72 }
73
74 protected override IEnumerable<T> GetValuesFromPage(ClientResult page)
75 {
76 using IEnumerator<T> enumerator = new SseUpdateEnumerator<T>(_eventDeserializerFunc, page, _cancellationToken, AdditionalDisposalActions);
77 while (enumerator.MoveNext())
78 {
79 yield return enumerator.Current;
80 }
81 }
82
83 private sealed class SseUpdateEnumerator<U> : IEnumerator<U>
84 {
85 private static ReadOnlySpan<byte> TerminalData => "[DONE]"u8;
86
87 private List<Action> _additionalDisposalActions;
88
89 private readonly CancellationToken _cancellationToken;
90 private readonly PipelineResponse _response;
91
92 // These enumerators represent what is effectively a doubly-nested
93 // loop over the outer event collection and the inner update collection,
94 // i.e.:
95 // foreach (var sse in _events) {
96 // // get _updates from sse event
97 // foreach (var update in _updates) { ... }
98 // }
99 private IEnumerator<SseItem<byte[]>>? _events;
100 private IEnumerator<U>? _updates;
101 private readonly Func<SseItem<byte[]>, IEnumerable<U>> _eventDeserializerFunc;
102
103 private U? _current;
104 private bool _started;
105
106 public SseUpdateEnumerator(
107 Func<SseItem<byte[]>, IEnumerable<U>> eventDeserializerFunc,
108 ClientResult page,
109 CancellationToken cancellationToken,
110 List<Action> additionalDisposalActions)
111 {
112 Argument.AssertNotNull(eventDeserializerFunc, nameof(eventDeserializerFunc));
113 Argument.AssertNotNull(page, nameof(page));
114
115 _eventDeserializerFunc = eventDeserializerFunc;
116 _response = page.GetRawResponse();
117 _cancellationToken = cancellationToken;
118 _additionalDisposalActions = additionalDisposalActions;
119 }
120
121 U IEnumerator<U>.Current => _current!;
122
123 object IEnumerator.Current => _current!;
124
125 public bool MoveNext()
126 {
127 if (_events is null && _started)
128 {
129 throw new ObjectDisposedException(typeof(U).Name);
130 }
131
132 _cancellationToken.ThrowIfCancellationRequested();
133 _events ??= CreateEventEnumerator();
134 _started = true;
135
136 if (_updates is not null && _updates.MoveNext())
137 {
138 _current = _updates.Current;
139 return true;
140 }
141
142 if (_events.MoveNext())
143 {
144 if (_events.Current.Data.AsSpan().SequenceEqual(TerminalData))
145 {
146 _current = default;
147 return false;
148 }
149
150 _updates = _eventDeserializerFunc.Invoke(_events.Current).GetEnumerator();
151
152 if (_updates.MoveNext())
153 {
154 _current = _updates.Current;
155 return true;
156 }
157 }
158
159 _current = default;
160 return false;
161 }
162
163 private IEnumerator<SseItem<byte[]>> CreateEventEnumerator()
164 {
165 if (_response.ContentStream is null)
166 {
167 throw new InvalidOperationException("Unable to create result from response with null ContentStream");
168 }
169
170 IEnumerable<SseItem<byte[]>> enumerable = SseParser.Create(_response.ContentStream, (_, bytes) => bytes.ToArray()).Enumerate();
171 return enumerable.GetEnumerator();
172 }
173
174 public void Reset()
175 {
176 throw new NotSupportedException("Cannot seek back in an SSE stream.");
177 }
178
179 public void Dispose()
180 {
181 Dispose(true);
182 GC.SuppressFinalize(this);
183 }
184
185 private void Dispose(bool disposing)
186 {
187 if (disposing && _events is not null)
188 {
189 _events.Dispose();
190 _events = null;
191
192 // Dispose the response so we don't leave the network connection open.
193 _response?.Dispose();
194 }
195
196 foreach (Action additionalDisposalAction in _additionalDisposalActions ?? [])
197 {
198 additionalDisposalAction?.Invoke();
199 }
200 _additionalDisposalActions?.Clear();
201 }
202 }
203}
204