@@ -185,33 +185,31 @@ public final class GenerativeModel: Sendable {
185
185
isStreaming: true ,
186
186
options: requestOptions)
187
187
188
- var responseIterator = generativeAIService. loadRequestStream ( request: generateContentRequest)
189
- . makeAsyncIterator ( )
190
- return AsyncThrowingStream {
191
- let response : GenerateContentResponse ?
192
- do {
193
- response = try await responseIterator. next ( )
194
- } catch {
195
- throw GenerativeModel . generateContentError ( from: error)
196
- }
197
-
198
- // The responseIterator will return `nil` when it's done.
199
- guard let response = response else {
200
- // This is the end of the stream! Signal it by sending `nil`.
201
- return nil
202
- }
188
+ return AsyncThrowingStream { continuation in
189
+ let responseStream = generativeAIService. loadRequestStream ( request: generateContentRequest)
190
+ Task {
191
+ do {
192
+ for try await response in responseStream {
193
+ // Check the prompt feedback to see if the prompt was blocked.
194
+ if response. promptFeedback? . blockReason != nil {
195
+ throw GenerateContentError . promptBlocked ( response: response)
196
+ }
203
197
204
- // Check the prompt feedback to see if the prompt was blocked.
205
- if response. promptFeedback? . blockReason != nil {
206
- throw GenerateContentError . promptBlocked ( response: response)
207
- }
198
+ // If the stream ended early unexpectedly, throw an error.
199
+ if let finishReason = response. candidates. first? . finishReason, finishReason != . stop {
200
+ throw GenerateContentError . responseStoppedEarly (
201
+ reason: finishReason,
202
+ response: response
203
+ )
204
+ }
208
205
209
- // If the stream ended early unexpectedly, throw an error.
210
- if let finishReason = response. candidates. first? . finishReason, finishReason != . stop {
211
- throw GenerateContentError . responseStoppedEarly ( reason: finishReason, response: response)
212
- } else {
213
- // Response was valid content, pass it along and continue.
214
- return response
206
+ continuation. yield ( response)
207
+ }
208
+ continuation. finish ( )
209
+ } catch {
210
+ continuation. finish ( throwing: GenerativeModel . generateContentError ( from: error) )
211
+ return
212
+ }
215
213
}
216
214
}
217
215
}
0 commit comments