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