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/StreamingUpdateCollection.cs

157lines · modecode

RawDownload
Latest commit unavailable.
1using System;
2using System.ClientModel;
3using System.ClientModel.Primitives;
4using System.Collections;
5using System.Collections.Generic;
6using System.Net.ServerSentEvents;
7using System.Threading;
8
9#nullable enable
10
11namespace OpenAI.Assistants;
12
13/// <summary>
14/// Implementation of collection abstraction over streaming assistant updates.
15/// </summary>
16internal class StreamingUpdateCollection : CollectionResult<StreamingUpdate>
17{
18 private readonly Func<ClientResult> _sendRequest;
19 private readonly CancellationToken _cancellationToken;
20
21 public StreamingUpdateCollection(
22 Func<ClientResult> sendRequest,
23 CancellationToken cancellationToken)
24 {
25 Argument.AssertNotNull(sendRequest, nameof(sendRequest));
26
27 _sendRequest = sendRequest;
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 override IEnumerable<ClientResult> GetRawPages()
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 _sendRequest();
40 }
41 protected override IEnumerable<StreamingUpdate> GetValuesFromPage(ClientResult page)
42 {
43 using IEnumerator<StreamingUpdate> enumerator = new StreamingUpdateEnumerator(page, _cancellationToken);
44 while (enumerator.MoveNext())
45 {
46 yield return enumerator.Current;
47 }
48 }
49
50 private sealed class StreamingUpdateEnumerator : IEnumerator<StreamingUpdate>
51 {
52 private static ReadOnlySpan<byte> TerminalData => "[DONE]"u8;
53
54 private readonly CancellationToken _cancellationToken;
55 private readonly PipelineResponse _response;
56
57 // These enumerators represent what is effectively a doubly-nested
58 // loop over the outer event collection and the inner update collection,
59 // i.e.:
60 // foreach (var sse in _events) {
61 // // get _updates from sse event
62 // foreach (var update in _updates) { ... }
63 // }
64 private IEnumerator<SseItem<byte[]>>? _events;
65 private IEnumerator<StreamingUpdate>? _updates;
66
67 private StreamingUpdate? _current;
68 private bool _started;
69
70 public StreamingUpdateEnumerator(ClientResult page, CancellationToken cancellationToken)
71 {
72 Argument.AssertNotNull(page, nameof(page));
73
74 _response = page.GetRawResponse();
75 _cancellationToken = cancellationToken;
76 }
77
78 StreamingUpdate IEnumerator<StreamingUpdate>.Current
79 => _current!;
80
81 object IEnumerator.Current => _current!;
82
83 public bool MoveNext()
84 {
85 if (_events is null && _started)
86 {
87 throw new ObjectDisposedException(nameof(StreamingUpdateEnumerator));
88 }
89
90
91 _cancellationToken.ThrowIfCancellationRequested();
92 _events ??= CreateEventEnumerator();
93 _started = true;
94
95 if (_updates is not null && _updates.MoveNext())
96 {
97 _current = _updates.Current;
98 return true;
99 }
100
101 if (_events.MoveNext())
102 {
103 if (_events.Current.Data.AsSpan().SequenceEqual(TerminalData))
104 {
105 _current = default;
106 return false;
107 }
108
109 var updates = StreamingUpdate.FromEvent(_events.Current);
110 _updates = updates.GetEnumerator();
111
112 if (_updates.MoveNext())
113 {
114 _current = _updates.Current;
115 return true;
116 }
117 }
118
119 _current = default;
120 return false;
121 }
122
123 private IEnumerator<SseItem<byte[]>> CreateEventEnumerator()
124 {
125 if (_response.ContentStream is null)
126 {
127 throw new InvalidOperationException("Unable to create result from response with null ContentStream");
128 }
129
130 IEnumerable<SseItem<byte[]>> enumerable = SseParser.Create(_response.ContentStream, (_, bytes) => bytes.ToArray()).Enumerate();
131 return enumerable.GetEnumerator();
132 }
133
134 public void Reset()
135 {
136 throw new NotSupportedException("Cannot seek back in an SSE stream.");
137 }
138
139 public void Dispose()
140 {
141 Dispose(true);
142 GC.SuppressFinalize(this);
143 }
144
145 private void Dispose(bool disposing)
146 {
147 if (disposing && _events is not null)
148 {
149 _events.Dispose();
150 _events = null;
151
152 // Dispose the response so we don't leave the network connection open.
153 _response?.Dispose();
154 }
155 }
156 }
157}
158