Remove uses of unstructured concurrency

andrewheard · andrewheard · commit c13d297a119e · 2024-08-27T13:02:07.000-04:00
diff --git a/FirebaseVertexAI/Sources/Chat.swift b/FirebaseVertexAI/Sources/Chat.swift
@@ -85,68 +85,62 @@ public actor Chat {
   /// - Parameter parts: The new content to send as a single chat message.
   /// - Returns: A stream containing the model's response or an error if an error occurred.
   @available(macOS 12.0, *)
-  public func sendMessageStream(_ parts: any ThrowingPartsRepresentable...)
+  public func sendMessageStream(_ parts: any ThrowingPartsRepresentable...) async throws
     -> AsyncThrowingStream<GenerateContentResponse, Error> {
-    return try sendMessageStream([ModelContent(parts: parts)])
+    return try await sendMessageStream([ModelContent(parts: parts)])
   }
 
   /// Sends a message using the existing history of this chat as context. If successful, the message
   /// and response will be added to the history. If unsuccessful, history will remain unchanged.
   /// - Parameter content: The new content to send as a single chat message.
   /// - Returns: A stream containing the model's response or an error if an error occurred.
   @available(macOS 12.0, *)
-  public func sendMessageStream(_ content: @autoclosure () throws -> [ModelContent])
+  public func sendMessageStream(_ content: @autoclosure () throws -> [ModelContent]) async throws
     -> AsyncThrowingStream<GenerateContentResponse, Error> {
     let resolvedContent: [ModelContent]
     do {
       resolvedContent = try content()
     } catch let underlying {
-      return AsyncThrowingStream { continuation in
-        let error: Error
+      // TODO: Consider throwing this before the stream.
+      return AsyncThrowingStream {
         if let contentError = underlying as? ImageConversionError {
-          error = GenerateContentError.promptImageContentError(underlying: contentError)
+          throw GenerateContentError.promptImageContentError(underlying: contentError)
         } else {
-          error = GenerateContentError.internalError(underlying: underlying)
+          throw GenerateContentError.internalError(underlying: underlying)
         }
-        continuation.finish(throwing: error)
       }
     }
 
-    return AsyncThrowingStream { continuation in
-      Task {
-        var aggregatedContent: [ModelContent] = []
-
-        // Ensure that the new content has the role set.
-        let newContent: [ModelContent] = resolvedContent.map(populateContentRole(_:))
-
-        // Send the history alongside the new message as context.
-        let request = history + newContent
-        let stream = await model.generateContentStream(request)
-        do {
-          for try await chunk in stream {
-            // Capture any content that's streaming. This should be populated if there's no error.
-            if let chunkContent = chunk.candidates.first?.content {
-              aggregatedContent.append(chunkContent)
-            }
-
-            // Pass along the chunk.
-            continuation.yield(chunk)
-          }
-        } catch {
-          // Rethrow the error that the underlying stream threw. Don't add anything to history.
-          continuation.finish(throwing: error)
-          return
-        }
+    var aggregatedContent: [ModelContent] = []
+
+    // Ensure that the new content has the role set.
+    let newContent: [ModelContent] = resolvedContent.map(populateContentRole(_:))
 
-        // Save the request.
-        history.append(contentsOf: newContent)
+    // Send the history alongside the new message as context.
+    let request = history + newContent
+    let stream = try await model.generateContentStream(request)
 
-        // Aggregate the content to add it to the history before we finish.
-        let aggregated = aggregatedChunks(aggregatedContent)
-        history.append(aggregated)
+    var streamIterator = stream.makeAsyncIterator()
 
-        continuation.finish()
+    return AsyncThrowingStream {
+      while let chunk = try await streamIterator.next() {
+        // Capture any content that's streaming. This should be populated if there's no error.
+        if let chunkContent = chunk.candidates.first?.content {
+          aggregatedContent.append(chunkContent)
+        }
+
+        // Pass along the chunk.
+        return chunk
       }
+
+      // Save the request.
+      self.history.append(contentsOf: newContent)
+
+      // Aggregate the content to add it to the history before we finish.
+      let aggregated = self.aggregatedChunks(aggregatedContent)
+      self.history.append(aggregated)
+
+      return nil
     }
   }
 
diff --git a/FirebaseVertexAI/Sources/GenerativeAIService.swift b/FirebaseVertexAI/Sources/GenerativeAIService.swift
@@ -72,94 +72,62 @@ struct GenerativeAIService {
   }
 
   @available(macOS 12.0, *)
-  func loadRequestStream<T: GenerativeAIRequest>(request: T)
+  func loadRequestStream<T: GenerativeAIRequest>(request: T) async throws
     -> AsyncThrowingStream<T.Response, Error> {
-    return AsyncThrowingStream { continuation in
-      Task {
-        let urlRequest: URLRequest
-        do {
-          urlRequest = try await self.urlRequest(request: request)
-        } catch {
-          continuation.finish(throwing: error)
-          return
-        }
+    let urlRequest = try await self.urlRequest(request: request)
 
-        #if DEBUG
-          printCURLCommand(from: urlRequest)
-        #endif
-
-        let stream: URLSession.AsyncBytes
-        let rawResponse: URLResponse
-        do {
-          (stream, rawResponse) = try await urlSession.bytes(for: urlRequest)
-        } catch {
-          continuation.finish(throwing: error)
-          return
-        }
+    #if DEBUG
+      printCURLCommand(from: urlRequest)
+    #endif
 
-        // Verify the status code is 200
-        let response: HTTPURLResponse
-        do {
-          response = try httpResponse(urlResponse: rawResponse)
-        } catch {
-          continuation.finish(throwing: error)
-          return
-        }
+    let stream: URLSession.AsyncBytes
+    let rawResponse: URLResponse
+    (stream, rawResponse) = try await urlSession.bytes(for: urlRequest)
+
+    let response = try httpResponse(urlResponse: rawResponse)
+
+    // Verify the status code is 200
+    guard response.statusCode == 200 else {
+      Logging.network
+        .error("[FirebaseVertexAI] The server responded with an error: \(response)")
+      var responseBody = ""
+      for try await line in stream.lines {
+        responseBody += line + "\n"
+      }
 
-        // Verify the status code is 200
-        guard response.statusCode == 200 else {
-          Logging.network
-            .error("[FirebaseVertexAI] The server responded with an error: \(response)")
-          var responseBody = ""
-          for try await line in stream.lines {
-            responseBody += line + "\n"
-          }
+      Logging.default.error("[FirebaseVertexAI] Response payload: \(responseBody)")
+      throw parseError(responseBody: responseBody)
+    }
 
-          Logging.default.error("[FirebaseVertexAI] Response payload: \(responseBody)")
-          continuation.finish(throwing: parseError(responseBody: responseBody))
+    // Received lines that are not server-sent events (SSE); these are not prefixed with "data:"
+    var extraLines = ""
 
-          return
-        }
+    let decoder = JSONDecoder()
+    decoder.keyDecodingStrategy = .convertFromSnakeCase
 
-        // Received lines that are not server-sent events (SSE); these are not prefixed with "data:"
-        var extraLines = ""
-
-        let decoder = JSONDecoder()
-        decoder.keyDecodingStrategy = .convertFromSnakeCase
-        for try await line in stream.lines {
-          Logging.network.debug("[FirebaseVertexAI] Stream response: \(line)")
-
-          if line.hasPrefix("data:") {
-            // We can assume 5 characters since it's utf-8 encoded, removing `data:`.
-            let jsonText = String(line.dropFirst(5))
-            let data: Data
-            do {
-              data = try jsonData(jsonText: jsonText)
-            } catch {
-              continuation.finish(throwing: error)
-              return
-            }
-
-            // Handle the content.
-            do {
-              let content = try parseResponse(T.Response.self, from: data)
-              continuation.yield(content)
-            } catch {
-              continuation.finish(throwing: error)
-              return
-            }
-          } else {
-            extraLines += line
-          }
-        }
+    var linesStream = stream.lines.makeAsyncIterator()
+
+    return AsyncThrowingStream<T.Response, Error> {
+      while let line = try await linesStream.next() {
+        Logging.network.debug("[FirebaseVertexAI] Stream response: \(line)")
+
+        if line.hasPrefix("data:") {
+          // We can assume 5 characters since it's utf-8 encoded, removing `data:`.
+          let jsonText = String(line.dropFirst(5))
+          let data = try jsonData(jsonText: jsonText)
 
-        if extraLines.count > 0 {
-          continuation.finish(throwing: parseError(responseBody: extraLines))
-          return
+          // Handle the content.
+          return try parseResponse(T.Response.self, from: data)
+        } else {
+          extraLines += line
         }
+      }
 
-        continuation.finish(throwing: nil)
+      if extraLines.count > 0 {
+        throw parseError(responseBody: extraLines)
       }
+
+      return nil
     }
   }
 
diff --git a/FirebaseVertexAI/Sources/GenerativeModel.swift b/FirebaseVertexAI/Sources/GenerativeModel.swift
@@ -179,9 +179,9 @@ public final actor GenerativeModel {
   /// - Returns: A stream wrapping content generated by the model or a ``GenerateContentError``
   ///     error if an error occurred.
   @available(macOS 12.0, *)
-  public func generateContentStream(_ parts: any ThrowingPartsRepresentable...)
+  public func generateContentStream(_ parts: any ThrowingPartsRepresentable...) async throws
     -> AsyncThrowingStream<GenerateContentResponse, Error> {
-    return try generateContentStream([ModelContent(parts: parts)])
+    return try await generateContentStream([ModelContent(parts: parts)])
   }
 
   /// Generates new content from input content given to the model as a prompt.
@@ -190,20 +190,19 @@ public final actor GenerativeModel {
   /// - Returns: A stream wrapping content generated by the model or a ``GenerateContentError``
   ///     error if an error occurred.
   @available(macOS 12.0, *)
-  public func generateContentStream(_ content: @autoclosure () throws -> [ModelContent])
+  public func generateContentStream(_ content: @autoclosure () throws
+    -> [ModelContent]) async throws
     -> AsyncThrowingStream<GenerateContentResponse, Error> {
     let evaluatedContent: [ModelContent]
     do {
       evaluatedContent = try content()
     } catch let underlying {
-      return AsyncThrowingStream { continuation in
-        let error: Error
+      // TODO: Consider throwing this before the stream.
+      return AsyncThrowingStream {
         if let contentError = underlying as? ImageConversionError {
-          error = GenerateContentError.promptImageContentError(underlying: contentError)
-        } else {
-          error = GenerateContentError.internalError(underlying: underlying)
+          throw GenerateContentError.promptImageContentError(underlying: contentError)
         }
-        continuation.finish(throwing: error)
+        throw GenerateContentError.internalError(underlying: underlying)
       }
     }
 
@@ -217,8 +216,15 @@ public final actor GenerativeModel {
                                                         isStreaming: true,
                                                         options: requestOptions)
 
-    var responseIterator = generativeAIService.loadRequestStream(request: generateContentRequest)
-      .makeAsyncIterator()
+    let responseStream: AsyncThrowingStream<GenerateContentResponse, Error>
+    do {
+      responseStream = try await generativeAIService
+        .loadRequestStream(request: generateContentRequest)
+    } catch {
+      throw GenerativeModel.generateContentError(from: error)
+    }
+
+    var responseIterator = responseStream.makeAsyncIterator()
     return AsyncThrowingStream {
       let response: GenerateContentResponse?
       do {
diff --git a/FirebaseVertexAI/Tests/Unit/ChatTests.swift b/FirebaseVertexAI/Tests/Unit/ChatTests.swift
@@ -64,7 +64,7 @@ final class ChatTests: XCTestCase {
     )
     let chat = Chat(model: model, history: [])
     let input = "Test input"
-    let stream = await chat.sendMessageStream(input)
+    let stream = try await chat.sendMessageStream(input)
 
     // Ensure the values are parsed correctly
     for try await value in stream {
diff --git a/FirebaseVertexAI/Tests/Unit/GenerativeModelTests.swift b/FirebaseVertexAI/Tests/Unit/GenerativeModelTests.swift

Original file line number	Diff line number	Diff line change
`@@ -64,7 +64,7 @@ final class ChatTests: XCTestCase {`
`64`	`64`	`)`
`65`	`65`	`let chat = Chat(model: model, history: [])`
`66`	`66`	`let input = "Test input"`
`67`		`- let stream = await chat.sendMessageStream(input)`
	`67`	`+ let stream = try await chat.sendMessageStream(input)`
`68`	`68`
`69`	`69`	`// Ensure the values are parsed correctly`
`70`	`70`	`for try await value in stream {`