openai/openai-dotnet

Public

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

CodeCommitsIssuesPull requestsActionsInsightsSecurity
OpenAI_2.0.0-beta.1

Branches

Tags

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

Clone

HTTPS

Download ZIP

src/Custom/Assistants/Streaming/StreamingUpdateCollection.cs

144lines · modecode

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