openai/openai-dotnet

Public

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

CodeCommitsIssuesPull requestsActionsInsightsSecurity
OpenAI_2.0.0-beta.7

Branches

Tags

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

Clone

HTTPS

Download ZIP

src/Custom/Chat/Internal/StreamingChatCompletionUpdateCollection.cs

147lines · modecode

1using System;
2using System.ClientModel;
3using System.ClientModel.Primitives;
4using System.Collections;
5using System.Collections.Generic;
6using System.Diagnostics;
7using System.Net.ServerSentEvents;
8using System.Text.Json;
9
10#nullable enable
11
12namespace OpenAI.Chat;
13
14/// <summary>
15/// Implementation of collection abstraction over streaming chat updates.
16/// </summary>
17internal class StreamingChatCompletionUpdateCollection : ResultCollection<StreamingChatCompletionUpdate>
18{
19 private readonly Func<ClientResult> _getResult;
20
21 public StreamingChatCompletionUpdateCollection(Func<ClientResult> getResult) : base()
22 {
23 Argument.AssertNotNull(getResult, nameof(getResult));
24
25 _getResult = getResult;
26 }
27
28 public override IEnumerator<StreamingChatCompletionUpdate> GetEnumerator()
29 {
30 return new StreamingChatUpdateEnumerator(_getResult, this);
31 }
32
33 private sealed class StreamingChatUpdateEnumerator : IEnumerator<StreamingChatCompletionUpdate>
34 {
35 private static ReadOnlySpan<byte> TerminalData => "[DONE]"u8;
36
37 private readonly Func<ClientResult> _getResult;
38 private readonly StreamingChatCompletionUpdateCollection _enumerable;
39
40 // These enumerators represent what is effectively a doubly-nested
41 // loop over the outer event collection and the inner update collection,
42 // i.e.:
43 // foreach (var sse in _events) {
44 // // get _updates from sse event
45 // foreach (var update in _updates) { ... }
46 // }
47 private IEnumerator<SseItem<byte[]>>? _events;
48 private IEnumerator<StreamingChatCompletionUpdate>? _updates;
49
50 private StreamingChatCompletionUpdate? _current;
51 private bool _started;
52
53 public StreamingChatUpdateEnumerator(Func<ClientResult> getResult,
54 StreamingChatCompletionUpdateCollection enumerable)
55 {
56 Debug.Assert(getResult is not null);
57 Debug.Assert(enumerable is not null);
58
59 _getResult = getResult!;
60 _enumerable = enumerable!;
61 }
62
63 StreamingChatCompletionUpdate IEnumerator<StreamingChatCompletionUpdate>.Current
64 => _current!;
65
66 object IEnumerator.Current => throw new NotImplementedException();
67
68 public bool MoveNext()
69 {
70 if (_events is null && _started)
71 {
72 throw new ObjectDisposedException(nameof(StreamingChatUpdateEnumerator));
73 }
74
75 _events ??= CreateEventEnumerator();
76 _started = true;
77
78 if (_updates is not null && _updates.MoveNext())
79 {
80 _current = _updates.Current;
81 return true;
82 }
83
84 if (_events.MoveNext())
85 {
86 if (_events.Current.Data.AsSpan().SequenceEqual(TerminalData))
87 {
88 _current = default;
89 return false;
90 }
91
92 using JsonDocument doc = JsonDocument.Parse(_events.Current.Data);
93 var updates = StreamingChatCompletionUpdate.DeserializeStreamingChatCompletionUpdates(doc.RootElement);
94 _updates = updates.GetEnumerator();
95
96 if (_updates.MoveNext())
97 {
98 _current = _updates.Current;
99 return true;
100 }
101 }
102
103 _current = default;
104 return false;
105 }
106
107 private IEnumerator<SseItem<byte[]>> CreateEventEnumerator()
108 {
109 ClientResult result = _getResult();
110 PipelineResponse response = result.GetRawResponse();
111 _enumerable.SetRawResponse(response);
112
113 if (response.ContentStream is null)
114 {
115 throw new InvalidOperationException("Unable to create result from response with null ContentStream");
116 }
117
118 IEnumerable<SseItem<byte[]>> enumerable = SseParser.Create(response.ContentStream, (_, bytes) => bytes.ToArray()).Enumerate();
119 return enumerable.GetEnumerator();
120 }
121
122 public void Reset()
123 {
124 throw new NotSupportedException("Cannot seek back in an SSE stream.");
125 }
126
127 public void Dispose()
128 {
129 Dispose(true);
130 GC.SuppressFinalize(this);
131 }
132
133 private void Dispose(bool disposing)
134 {
135 if (disposing && _events is not null)
136 {
137 _events.Dispose();
138 _events = null;
139
140 // Dispose the response so we don't leave the unbuffered
141 // network stream open.
142 PipelineResponse response = _enumerable.GetRawResponse();
143 response.Dispose();
144 }
145 }
146 }
147}
148