openai/openai-dotnet

Public

mirrored from https://github.com/openai/openai-dotnetAvailable

CodeCommitsIssuesPull requestsActionsInsightsSecurity
OpenAI_2.1.0

Branches

Tags

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

Clone

HTTPS

Download ZIP

src/Custom/Chat/Streaming/InternalStreamingChatCompletionUpdateCollection.cs

159lines · modecode

1using System;
2using System.ClientModel;
3using System.ClientModel.Primitives;
4using System.Collections;
5using System.Collections.Generic;
6using System.Net.ServerSentEvents;
7using System.Text.Json;
8using System.Threading;
9
10#nullable enable
11
12namespace OpenAI.Chat;
13
14/// <summary>
15/// Implementation of collection abstraction over streaming chat updates.
16/// </summary>
17internal class InternalStreamingChatCompletionUpdateCollection : CollectionResult<StreamingChatCompletionUpdate>
18{
19 private readonly Func<ClientResult> _sendRequest;
20 private readonly CancellationToken _cancellationToken;
21
22 public InternalStreamingChatCompletionUpdateCollection(
23 Func<ClientResult> sendRequest,
24 CancellationToken cancellationToken)
25 {
26 Argument.AssertNotNull(sendRequest, nameof(sendRequest));
27
28 _sendRequest = sendRequest;
29 _cancellationToken = cancellationToken;
30 }
31
32 public override ContinuationToken? GetContinuationToken(ClientResult page)
33 // Continuation is not supported for SSE streams.
34 => null;
35
36 public override IEnumerable<ClientResult> GetRawPages()
37 {
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 }
50 }
51
52 private sealed class StreamingChatUpdateEnumerator : IEnumerator<StreamingChatCompletionUpdate>
53 {
54 private static ReadOnlySpan<byte> TerminalData => "[DONE]"u8;
55
56 private readonly CancellationToken _cancellationToken;
57 private readonly PipelineResponse _response;
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 // }
66 private IEnumerator<SseItem<byte[]>>? _events;
67 private IEnumerator<StreamingChatCompletionUpdate>? _updates;
68
69 private StreamingChatCompletionUpdate? _current;
70 private bool _started;
71
72 public StreamingChatUpdateEnumerator(ClientResult page, CancellationToken cancellationToken)
73 {
74 Argument.AssertNotNull(page, nameof(page));
75
76 _response = page.GetRawResponse();
77 _cancellationToken = cancellationToken;
78 }
79
80 StreamingChatCompletionUpdate IEnumerator<StreamingChatCompletionUpdate>.Current
81 => _current!;
82
83 object IEnumerator.Current => _current!;
84
85 public bool MoveNext()
86 {
87 if (_events is null && _started)
88 {
89 throw new ObjectDisposedException(nameof(StreamingChatUpdateEnumerator));
90 }
91
92 _cancellationToken.ThrowIfCancellationRequested();
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 {
104 if (_events.Current.Data.AsSpan().SequenceEqual(TerminalData))
105 {
106 _current = default;
107 return false;
108 }
109
110 using JsonDocument doc = JsonDocument.Parse(_events.Current.Data);
111 List<StreamingChatCompletionUpdate> updates = [StreamingChatCompletionUpdate.DeserializeStreamingChatCompletionUpdate(doc.RootElement)];
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
125 private IEnumerator<SseItem<byte[]>> CreateEventEnumerator()
126 {
127 if (_response.ContentStream is null)
128 {
129 throw new InvalidOperationException("Unable to create result from response with null ContentStream");
130 }
131
132 IEnumerable<SseItem<byte[]>> enumerable = SseParser.Create(_response.ContentStream, (_, bytes) => bytes.ToArray()).Enumerate();
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
154 // Dispose the response so we don't leave the network connection open.
155 _response?.Dispose();
156 }
157 }
158 }
159}
160