Make GenerativeModel and Chat into Swift actors

Author: andrewheard

FirebaseVertexAI/Sources/Chat.swift

@@ -17,7 +17,7 @@ import Foundation
 /// An object that represents a back-and-forth chat with a model, capturing the history and saving
 /// the context in memory between each message sent.
 @available(iOS 15.0, macOS 11.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
-public class Chat {
+public actor Chat {
   private let model: GenerativeModel
 
   /// Initializes a new chat representing a 1:1 conversation between model and user.
@@ -121,7 +121,7 @@ public class Chat {
 
         // Send the history alongside the new message as context.
         let request = history + newContent
-        let stream = model.generateContentStream(request)
+        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.

FirebaseVertexAI/Sources/GenerativeModel.swift

@@ -19,7 +19,7 @@ import Foundation
 /// A type that represents a remote multimodal model (like Gemini), with the ability to generate
 /// content based on various input types.
 @available(iOS 15.0, macOS 11.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
-public final class GenerativeModel {
+public final actor GenerativeModel {
   /// The resource name of the model in the backend; has the format "models/model-name".
   let modelResourceName: String
 
@@ -219,31 +219,37 @@ public final class GenerativeModel {
 
     var responseIterator = generativeAIService.loadRequestStream(request: generateContentRequest)
       .makeAsyncIterator()
-    return AsyncThrowingStream {
-      let response: GenerateContentResponse?
-      do {
-        response = try await responseIterator.next()
-      } catch {
-        throw GenerativeModel.generateContentError(from: error)
-      }
+    return AsyncThrowingStream { contination in
+      Task {
+        do {
+          // The responseIterator will return `nil` when it's done.
+          while let response = try await responseIterator.next() {
+            // Check the prompt feedback to see if the prompt was blocked.
+            if response.promptFeedback?.blockReason != nil {
+              contination.finish(throwing: GenerateContentError.promptBlocked(response: response))
+              return
+            }
 
-      // The responseIterator will return `nil` when it's done.
-      guard let response = response else {
-        // This is the end of the stream! Signal it by sending `nil`.
-        return nil
-      }
+            // If the stream ended early unexpectedly, throw an error.
+            if let finishReason = response.candidates.first?.finishReason, finishReason != .stop {
+              contination.finish(throwing: GenerateContentError.responseStoppedEarly(
+                reason: finishReason,
+                response: response
+              ))
+              return
+            }
 
-      // Check the prompt feedback to see if the prompt was blocked.
-      if response.promptFeedback?.blockReason != nil {
-        throw GenerateContentError.promptBlocked(response: response)
-      }
+            // Response was valid content, pass it along and continue.
+            contination.yield(response)
+          }
 
-      // If the stream ended early unexpectedly, throw an error.
-      if let finishReason = response.candidates.first?.finishReason, finishReason != .stop {
-        throw GenerateContentError.responseStoppedEarly(reason: finishReason, response: response)
-      } else {
-        // Response was valid content, pass it along and continue.
-        return response
+          // This is the end of the stream! Signal it by calling `finish`.
+          contination.finish()
+          return
+        } catch {
+          contination.finish(throwing: GenerativeModel.generateContentError(from: error))
+          return
+        }
       }
     }
   }

FirebaseVertexAI/Tests/Unit/ChatTests.swift

@@ -64,19 +64,20 @@ final class ChatTests: XCTestCase {
     )
     let chat = Chat(model: model, history: [])
     let input = "Test input"
-    let stream = chat.sendMessageStream(input)
+    let stream = await chat.sendMessageStream(input)
 
     // Ensure the values are parsed correctly
     for try await value in stream {
       XCTAssertNotNil(value.text)
     }
 
-    XCTAssertEqual(chat.history.count, 2)
-    XCTAssertEqual(chat.history[0].parts[0].text, input)
+    let history = await chat.history
+    XCTAssertEqual(history.count, 2)
+    XCTAssertEqual(history[0].parts[0].text, input)
 
     let finalText = "1 2 3 4 5 6 7 8"
     let assembledExpectation = ModelContent(role: "model", parts: finalText)
-    XCTAssertEqual(chat.history[0].parts[0].text, input)
-    XCTAssertEqual(chat.history[1], assembledExpectation)
+    XCTAssertEqual(history[0].parts[0].text, input)
+    XCTAssertEqual(history[1], assembledExpectation)
   }
 }

FirebaseVertexAI/Tests/Unit/GenerativeModelTests.swift

@@ -760,7 +760,7 @@ final class GenerativeModelTests: XCTestCase {
       )
 
     do {
-      let stream = model.generateContentStream("Hi")
+      let stream = await model.generateContentStream("Hi")
       for try await _ in stream {
         XCTFail("No content is there, this shouldn't happen.")
       }
@@ -784,7 +784,7 @@ final class GenerativeModelTests: XCTestCase {
       )
 
     do {
-      let stream = model.generateContentStream(testPrompt)
+      let stream = await model.generateContentStream(testPrompt)
       for try await _ in stream {
         XCTFail("No content is there, this shouldn't happen.")
       }
@@ -807,7 +807,7 @@ final class GenerativeModelTests: XCTestCase {
       )
 
     do {
-      let stream = model.generateContentStream("Hi")
+      let stream = await model.generateContentStream("Hi")
       for try await _ in stream {
         XCTFail("No content is there, this shouldn't happen.")
       }
@@ -827,7 +827,7 @@ final class GenerativeModelTests: XCTestCase {
       )
 
     do {
-      let stream = model.generateContentStream("Hi")
+      let stream = await model.generateContentStream("Hi")
       for try await _ in stream {
         XCTFail("Content shouldn't be shown, this shouldn't happen.")
       }
@@ -847,7 +847,7 @@ final class GenerativeModelTests: XCTestCase {
       )
 
     do {
-      let stream = model.generateContentStream("Hi")
+      let stream = await model.generateContentStream("Hi")
       for try await _ in stream {
         XCTFail("Content shouldn't be shown, this shouldn't happen.")
       }
@@ -866,7 +866,7 @@ final class GenerativeModelTests: XCTestCase {
         withExtension: "txt"
       )
 
-    let stream = model.generateContentStream("Hi")
+    let stream = await model.generateContentStream("Hi")
     do {
       for try await content in stream {
         XCTAssertNotNil(content.text)
@@ -887,7 +887,7 @@ final class GenerativeModelTests: XCTestCase {
       )
 
     var responses = 0
-    let stream = model.generateContentStream("Hi")
+    let stream = await model.generateContentStream("Hi")
     for try await content in stream {
       XCTAssertNotNil(content.text)
       responses += 1
@@ -904,7 +904,7 @@ final class GenerativeModelTests: XCTestCase {
       )
 
     var responses = 0
-    let stream = model.generateContentStream("Hi")
+    let stream = await model.generateContentStream("Hi")
     for try await content in stream {
       XCTAssertNotNil(content.text)
       responses += 1
@@ -921,7 +921,7 @@ final class GenerativeModelTests: XCTestCase {
       )
 
     var hadUnknown = false
-    let stream = model.generateContentStream("Hi")
+    let stream = await model.generateContentStream("Hi")
     for try await content in stream {
       XCTAssertNotNil(content.text)
       if let ratings = content.candidates.first?.safetyRatings,
@@ -940,7 +940,7 @@ final class GenerativeModelTests: XCTestCase {
         withExtension: "txt"
       )
 
-    let stream = model.generateContentStream("Hi")
+    let stream = await model.generateContentStream("Hi")
     var citations = [Citation]()
     var responses = [GenerateContentResponse]()
     for try await content in stream {
@@ -996,7 +996,7 @@ final class GenerativeModelTests: XCTestCase {
         appCheckToken: appCheckToken
       )
 
-    let stream = model.generateContentStream(testPrompt)
+    let stream = await model.generateContentStream(testPrompt)
     for try await _ in stream {}
   }
 
@@ -1018,7 +1018,7 @@ final class GenerativeModelTests: XCTestCase {
         appCheckToken: AppCheckInteropFake.placeholderTokenValue
       )
 
-    let stream = model.generateContentStream(testPrompt)
+    let stream = await model.generateContentStream(testPrompt)
     for try await _ in stream {}
   }
 
@@ -1030,7 +1030,7 @@ final class GenerativeModelTests: XCTestCase {
       )
     var responses = [GenerateContentResponse]()
 
-    let stream = model.generateContentStream(testPrompt)
+    let stream = await model.generateContentStream(testPrompt)
     for try await response in stream {
       responses.append(response)
     }
@@ -1056,7 +1056,7 @@ final class GenerativeModelTests: XCTestCase {
 
     var responseCount = 0
     do {
-      let stream = model.generateContentStream("Hi")
+      let stream = await model.generateContentStream("Hi")
       for try await content in stream {
         XCTAssertNotNil(content.text)
         responseCount += 1
@@ -1076,7 +1076,7 @@ final class GenerativeModelTests: XCTestCase {
   func testGenerateContentStream_nonHTTPResponse() async throws {
     MockURLProtocol.requestHandler = try nonHTTPRequestHandler()
 
-    let stream = model.generateContentStream("Hi")
+    let stream = await model.generateContentStream("Hi")
     do {
       for try await content in stream {
         XCTFail("Unexpected content in stream: \(content)")
@@ -1096,7 +1096,7 @@ final class GenerativeModelTests: XCTestCase {
         withExtension: "txt"
       )
 
-    let stream = model.generateContentStream(testPrompt)
+    let stream = await model.generateContentStream(testPrompt)
     do {
       for try await content in stream {
         XCTFail("Unexpected content in stream: \(content)")
@@ -1120,7 +1120,7 @@ final class GenerativeModelTests: XCTestCase {
         withExtension: "txt"
       )
 
-    let stream = model.generateContentStream(testPrompt)
+    let stream = await model.generateContentStream(testPrompt)
     do {
       for try await content in stream {
         XCTFail("Unexpected content in stream: \(content)")
@@ -1159,7 +1159,7 @@ final class GenerativeModelTests: XCTestCase {
     )
 
     var responses = 0
-    let stream = model.generateContentStream(testPrompt)
+    let stream = await model.generateContentStream(testPrompt)
     for try await content in stream {
       XCTAssertNotNil(content.text)
       responses += 1

FirebaseVertexAI/Tests/Unit/VertexAIAPITests.swift

@@ -170,8 +170,8 @@ final class VertexAIAPITests: XCTestCase {
     #endif
 
     // Chat
-    _ = genAI.startChat()
-    _ = genAI.startChat(history: [ModelContent(parts: "abc")])
+    _ = await genAI.startChat()
+    _ = await genAI.startChat(history: [ModelContent(parts: "abc")])
   }
 
   // Public API tests for GenerateContentResponse.

FirebaseVertexAI/Tests/Unit/VertexComponentTests.swift

@@ -106,15 +106,20 @@ class VertexComponentTests: XCTestCase {
     let app = try XCTUnwrap(VertexComponentTests.app)
     let vertex = VertexAI.vertexAI(app: app, location: location)
     let modelName = "test-model-name"
-    let modelResourceName = vertex.modelResourceName(modelName: modelName)
-    let systemInstruction = ModelContent(role: "system", parts: "test-system-instruction-prompt")
+    let expectedModelResourceName = vertex.modelResourceName(modelName: modelName)
+    let expectedSystemInstruction = ModelContent(
+      role: "system",
+      parts: "test-system-instruction-prompt"
+    )
 
     let generativeModel = vertex.generativeModel(
       modelName: modelName,
-      systemInstruction: systemInstruction
+      systemInstruction: expectedSystemInstruction
     )
 
-    XCTAssertEqual(generativeModel.modelResourceName, modelResourceName)
-    XCTAssertEqual(generativeModel.systemInstruction, systemInstruction)
+    let modelResourceName = await generativeModel.modelResourceName
+    let systemInstruction = await generativeModel.systemInstruction
+    XCTAssertEqual(modelResourceName, expectedModelResourceName)
+    XCTAssertEqual(systemInstruction, expectedSystemInstruction)
   }
 }