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.1

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

143lines · modecode

RawDownload
Latest commit unavailable.
1using System;
2using System.ClientModel;
3using System.ClientModel.Primitives;
4using System.Collections.Generic;
5using System.Diagnostics;
6using System.Threading;
7using System.Threading.Tasks;
8
9#nullable enable
10
11namespace OpenAI.Assistants;
12
13/// <summary>
14/// Implementation of collection abstraction over streaming assistant updates.
15/// </summary>
16internal class AsyncStreamingUpdateCollection : AsyncResultCollection<StreamingUpdate>
17{
18 private readonly Func<Task<ClientResult>> _getResultAsync;
19
20 public AsyncStreamingUpdateCollection(Func<Task<ClientResult>> getResultAsync) : base()
21 {
22 Argument.AssertNotNull(getResultAsync, nameof(getResultAsync));
23
24 _getResultAsync = getResultAsync;
25 }
26
27 public override IAsyncEnumerator<StreamingUpdate> GetAsyncEnumerator(CancellationToken cancellationToken = default)
28 {
29 return new AsyncStreamingUpdateEnumerator(_getResultAsync, this, cancellationToken);
30 }
31
32 private sealed class AsyncStreamingUpdateEnumerator : IAsyncEnumerator<StreamingUpdate>
33 {
34 private const string _terminalData = "[DONE]";
35
36 private readonly Func<Task<ClientResult>> _getResultAsync;
37 private readonly AsyncStreamingUpdateCollection _enumerable;
38 private readonly CancellationToken _cancellationToken;
39
40 // These enumerators represent what is effectively a doubly-nested
41 // loop over the outer event collection and the inner update collection,
42 // i.e.:
43 // foreach (var sse in _events) {
44 // // get _updates from sse event
45 // foreach (var update in _updates) { ... }
46 // }
47 private IAsyncEnumerator<ServerSentEvent>? _events;
48 private IEnumerator<StreamingUpdate>? _updates;
49
50 private StreamingUpdate? _current;
51 private bool _started;
52
53 public AsyncStreamingUpdateEnumerator(Func<Task<ClientResult>> getResultAsync,
54 AsyncStreamingUpdateCollection enumerable,
55 CancellationToken cancellationToken)
56 {
57 Debug.Assert(getResultAsync is not null);
58 Debug.Assert(enumerable is not null);
59
60 _getResultAsync = getResultAsync!;
61 _enumerable = enumerable!;
62 _cancellationToken = cancellationToken;
63 }
64
65 StreamingUpdate IAsyncEnumerator<StreamingUpdate>.Current
66 => _current!;
67
68 async ValueTask<bool> IAsyncEnumerator<StreamingUpdate>.MoveNextAsync()
69 {
70 if (_events is null && _started)
71 {
72 throw new ObjectDisposedException(nameof(AsyncStreamingUpdateEnumerator));
73 }
74
75 _cancellationToken.ThrowIfCancellationRequested();
76 _events ??= await CreateEventEnumeratorAsync().ConfigureAwait(false);
77 _started = true;
78
79 if (_updates is not null && _updates.MoveNext())
80 {
81 _current = _updates.Current;
82 return true;
83 }
84
85 if (await _events.MoveNextAsync().ConfigureAwait(false))
86 {
87 if (_events.Current.Data == _terminalData)
88 {
89 _current = default;
90 return false;
91 }
92
93 var updates = StreamingUpdate.FromEvent(_events.Current);
94 _updates = updates.GetEnumerator();
95
96 if (_updates.MoveNext())
97 {
98 _current = _updates.Current;
99 return true;
100 }
101 }
102
103 _current = default;
104 return false;
105 }
106
107 private async Task<IAsyncEnumerator<ServerSentEvent>> CreateEventEnumeratorAsync()
108 {
109 ClientResult result = await _getResultAsync().ConfigureAwait(false);
110 PipelineResponse response = result.GetRawResponse();
111 _enumerable.SetRawResponse(response);
112
113 if (response.ContentStream is null)
114 {
115 throw new InvalidOperationException("Unable to create result from response with null ContentStream");
116 }
117
118 AsyncServerSentEventEnumerable enumerable = new(response.ContentStream);
119 return enumerable.GetAsyncEnumerator(_cancellationToken);
120 }
121
122 public async ValueTask DisposeAsync()
123 {
124 await DisposeAsyncCore().ConfigureAwait(false);
125
126 GC.SuppressFinalize(this);
127 }
128
129 private async ValueTask DisposeAsyncCore()
130 {
131 if (_events is not null)
132 {
133 await _events.DisposeAsync().ConfigureAwait(false);
134 _events = null;
135
136 // Dispose the response so we don't leave the unbuffered
137 // network stream open.
138 PipelineResponse response = _enumerable.GetRawResponse();
139 response.Dispose();
140 }
141 }
142 }
143}
144