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 · modeblame

58f93c8dShivangiReja2 years ago1using System;
9f9f2936Jose Arriaga Maldonado2 years ago2using System.ClientModel;
3using System.ClientModel.Primitives;
4using System.Collections;
5using System.Collections.Generic;
674e0f77Stephen Toub2 years ago6using System.Net.ServerSentEvents;
9f9f2936Jose Arriaga Maldonado2 years ago7using System.Text.Json;
2ab1a942Jose Arriaga Maldonado1 years ago8using System.Threading;
9f9f2936Jose Arriaga Maldonado2 years ago9
10#nullable enable
11
12namespace OpenAI.Chat;
13
14/// <summary>
15/// Implementation of collection abstraction over streaming chat updates.
16/// </summary>
2ab1a942Jose Arriaga Maldonado1 years ago17internal class InternalStreamingChatCompletionUpdateCollection : CollectionResult<StreamingChatCompletionUpdate>
9f9f2936Jose Arriaga Maldonado2 years ago18{
2ab1a942Jose Arriaga Maldonado1 years ago19private readonly Func<ClientResult> _sendRequest;
20private readonly CancellationToken _cancellationToken;
9f9f2936Jose Arriaga Maldonado2 years ago21
2ab1a942Jose Arriaga Maldonado1 years ago22public InternalStreamingChatCompletionUpdateCollection(
b0f9e5c3Jose Arriaga Maldonado1 years ago23Func<ClientResult> sendRequest,
2ab1a942Jose Arriaga Maldonado1 years ago24CancellationToken cancellationToken)
9f9f2936Jose Arriaga Maldonado2 years ago25{
2ab1a942Jose Arriaga Maldonado1 years ago26Argument.AssertNotNull(sendRequest, nameof(sendRequest));
9f9f2936Jose Arriaga Maldonado2 years ago27
2ab1a942Jose Arriaga Maldonado1 years ago28_sendRequest = sendRequest;
29_cancellationToken = cancellationToken;
9f9f2936Jose Arriaga Maldonado2 years ago30}
31
2ab1a942Jose Arriaga Maldonado1 years ago32public override ContinuationToken? GetContinuationToken(ClientResult page)
33// Continuation is not supported for SSE streams.
34=> null;
35
36public override IEnumerable<ClientResult> GetRawPages()
9f9f2936Jose Arriaga Maldonado2 years ago37{
2ab1a942Jose Arriaga Maldonado1 years ago38// We don't currently support resuming a dropped connection from the
39// last received event, so the response collection has a single element.
40yield return _sendRequest();
41}
42
43protected override IEnumerable<StreamingChatCompletionUpdate> GetValuesFromPage(ClientResult page)
44{
45using IEnumerator<StreamingChatCompletionUpdate> enumerator = new StreamingChatUpdateEnumerator(page, _cancellationToken);
46while (enumerator.MoveNext())
47{
48yield return enumerator.Current;
49}
9f9f2936Jose Arriaga Maldonado2 years ago50}
51
52private sealed class StreamingChatUpdateEnumerator : IEnumerator<StreamingChatCompletionUpdate>
53{
674e0f77Stephen Toub2 years ago54private static ReadOnlySpan<byte> TerminalData => "[DONE]"u8;
9f9f2936Jose Arriaga Maldonado2 years ago55
2ab1a942Jose Arriaga Maldonado1 years ago56private readonly CancellationToken _cancellationToken;
57private readonly PipelineResponse _response;
9f9f2936Jose Arriaga Maldonado2 years ago58
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 ago66private IEnumerator<SseItem<byte[]>>? _events;
9f9f2936Jose Arriaga Maldonado2 years ago67private IEnumerator<StreamingChatCompletionUpdate>? _updates;
68
69private StreamingChatCompletionUpdate? _current;
70private bool _started;
71
2ab1a942Jose Arriaga Maldonado1 years ago72public StreamingChatUpdateEnumerator(ClientResult page, CancellationToken cancellationToken)
9f9f2936Jose Arriaga Maldonado2 years ago73{
2ab1a942Jose Arriaga Maldonado1 years ago74Argument.AssertNotNull(page, nameof(page));
9f9f2936Jose Arriaga Maldonado2 years ago75
2ab1a942Jose Arriaga Maldonado1 years ago76_response = page.GetRawResponse();
77_cancellationToken = cancellationToken;
9f9f2936Jose Arriaga Maldonado2 years ago78}
79
80StreamingChatCompletionUpdate IEnumerator<StreamingChatCompletionUpdate>.Current
81=> _current!;
82
2ab1a942Jose Arriaga Maldonado1 years ago83object IEnumerator.Current => _current!;
9f9f2936Jose Arriaga Maldonado2 years ago84
85public bool MoveNext()
86{
87if (_events is null && _started)
88{
89throw new ObjectDisposedException(nameof(StreamingChatUpdateEnumerator));
90}
91
2ab1a942Jose Arriaga Maldonado1 years ago92_cancellationToken.ThrowIfCancellationRequested();
9f9f2936Jose Arriaga Maldonado2 years ago93_events ??= CreateEventEnumerator();
94_started = true;
95
96if (_updates is not null && _updates.MoveNext())
97{
98_current = _updates.Current;
99return true;
100}
101
102if (_events.MoveNext())
103{
674e0f77Stephen Toub2 years ago104if (_events.Current.Data.AsSpan().SequenceEqual(TerminalData))
9f9f2936Jose Arriaga Maldonado2 years ago105{
106_current = default;
107return false;
108}
109
110using JsonDocument doc = JsonDocument.Parse(_events.Current.Data);
31c2ba63Jose Arriaga Maldonado1 years ago111List<StreamingChatCompletionUpdate> updates = [StreamingChatCompletionUpdate.DeserializeStreamingChatCompletionUpdate(doc.RootElement)];
9f9f2936Jose Arriaga Maldonado2 years ago112_updates = updates.GetEnumerator();
113
114if (_updates.MoveNext())
115{
116_current = _updates.Current;
117return true;
118}
119}
120
121_current = default;
122return false;
123}
124
674e0f77Stephen Toub2 years ago125private IEnumerator<SseItem<byte[]>> CreateEventEnumerator()
9f9f2936Jose Arriaga Maldonado2 years ago126{
2ab1a942Jose Arriaga Maldonado1 years ago127if (_response.ContentStream is null)
9f9f2936Jose Arriaga Maldonado2 years ago128{
129throw new InvalidOperationException("Unable to create result from response with null ContentStream");
130}
131
2ab1a942Jose Arriaga Maldonado1 years ago132IEnumerable<SseItem<byte[]>> enumerable = SseParser.Create(_response.ContentStream, (_, bytes) => bytes.ToArray()).Enumerate();
9f9f2936Jose Arriaga Maldonado2 years ago133return enumerable.GetEnumerator();
134}
135
136public void Reset()
137{
138throw new NotSupportedException("Cannot seek back in an SSE stream.");
139}
140
141public void Dispose()
142{
143Dispose(true);
144GC.SuppressFinalize(this);
145}
146
147private void Dispose(bool disposing)
148{
149if (disposing && _events is not null)
150{
151_events.Dispose();
152_events = null;
153
2ab1a942Jose Arriaga Maldonado1 years ago154// Dispose the response so we don't leave the network connection open.
155_response?.Dispose();
9f9f2936Jose Arriaga Maldonado2 years ago156}
157}
158}
159}