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
674e0f773b26a22eb039e879539c4c7a44fdffdd

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

145lines · modecode

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