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

Branches

Tags

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

Clone

HTTPS

Download ZIP

openai-dotnet/src/Custom/Chat/Internal

src/Custom/Chat/Internal/AsyncStreamingChatCompletionUpdateCollection.cs

146lines · modecode

RawDownload
Latest commit unavailable.
1using System;
2using System.ClientModel;
3using System.ClientModel.Primitives;
4using System.Collections.Generic;
5using System.Diagnostics;
6using System.Net.ServerSentEvents;
7using System.Text.Json;
8using System.Threading;
9using System.Threading.Tasks;
10
11#nullable enable
12
13namespace OpenAI.Chat;
14
15/// <summary>
16/// Implementation of collection abstraction over streaming chat updates.
17/// </summary>
18internal class AsyncStreamingChatCompletionUpdateCollection : AsyncResultCollection<StreamingChatCompletionUpdate>
19{
20 private readonly Func<Task<ClientResult>> _getResultAsync;
21
22 public AsyncStreamingChatCompletionUpdateCollection(Func<Task<ClientResult>> getResultAsync) : base()
23 {
24 Argument.AssertNotNull(getResultAsync, nameof(getResultAsync));
25
26 _getResultAsync = getResultAsync;
27 }
28
29 public override IAsyncEnumerator<StreamingChatCompletionUpdate> GetAsyncEnumerator(CancellationToken cancellationToken = default)
30 {
31 return new AsyncStreamingChatUpdateEnumerator(_getResultAsync, this, cancellationToken);
32 }
33
34 private sealed class AsyncStreamingChatUpdateEnumerator : IAsyncEnumerator<StreamingChatCompletionUpdate>
35 {
36 private static ReadOnlySpan<byte> TerminalData => "[DONE]"u8;
37
38 private readonly Func<Task<ClientResult>> _getResultAsync;
39 private readonly AsyncStreamingChatCompletionUpdateCollection _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<StreamingChatCompletionUpdate>? _updates;
51
52 private StreamingChatCompletionUpdate? _current;
53 private bool _started;
54
55 public AsyncStreamingChatUpdateEnumerator(Func<Task<ClientResult>> getResultAsync,
56 AsyncStreamingChatCompletionUpdateCollection 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 StreamingChatCompletionUpdate IAsyncEnumerator<StreamingChatCompletionUpdate>.Current
68 => _current!;
69
70 async ValueTask<bool> IAsyncEnumerator<StreamingChatCompletionUpdate>.MoveNextAsync()
71 {
72 if (_events is null && _started)
73 {
74 throw new ObjectDisposedException(nameof(AsyncStreamingChatUpdateEnumerator));
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 using JsonDocument doc = JsonDocument.Parse(_events.Current.Data);
96 var updates = StreamingChatCompletionUpdate.DeserializeStreamingChatCompletionUpdates(doc.RootElement);
97 _updates = updates.GetEnumerator();
98
99 if (_updates.MoveNext())
100 {
101 _current = _updates.Current;
102 return true;
103 }
104 }
105
106 _current = default;
107 return false;
108 }
109
110 private async Task<IAsyncEnumerator<SseItem<byte[]>>> CreateEventEnumeratorAsync()
111 {
112 ClientResult result = await _getResultAsync().ConfigureAwait(false);
113 PipelineResponse response = result.GetRawResponse();
114 _enumerable.SetRawResponse(response);
115
116 if (response.ContentStream is null)
117 {
118 throw new InvalidOperationException("Unable to create result from response with null ContentStream");
119 }
120
121 IAsyncEnumerable<SseItem<byte[]>> enumerable = SseParser.Create(response.ContentStream, (_, bytes) => bytes.ToArray()).EnumerateAsync();
122 return enumerable.GetAsyncEnumerator(_cancellationToken);
123 }
124
125 public async ValueTask DisposeAsync()
126 {
127 await DisposeAsyncCore().ConfigureAwait(false);
128
129 GC.SuppressFinalize(this);
130 }
131
132 private async ValueTask DisposeAsyncCore()
133 {
134 if (_events is not null)
135 {
136 await _events.DisposeAsync().ConfigureAwait(false);
137 _events = null;
138
139 // Dispose the response so we don't leave the unbuffered
140 // network stream open.
141 PipelineResponse response = _enumerable.GetRawResponse();
142 response.Dispose();
143 }
144 }
145 }
146}