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

Branches

Tags

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

Clone

HTTPS

Download ZIP

openai-dotnet/src/Custom/Chat/Streaming

src/Custom/Chat/Streaming/InternalAsyncStreamingChatCompletionUpdateCollection.cs

154lines · modecode

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