[Vertex AI] Make GenerativeModel and Chat into Swift actors (#13545)

Author: andrewheard

FirebaseVertexAI/CHANGELOG.md

@@ -1,6 +1,9 @@
 # 11.2.0
 - [fixed] Resolved a decoding error for citations without a `uri` and added
   support for decoding `title` fields, which were previously ignored. (#13518)
+- [changed] **Breaking Change**: The methods for starting streaming requests
+  (`generateContentStream` and `sendMessageStream`) and creating a chat instance
+  (`startChat`) are now asynchronous and must be called with `await`. (#13545)
 
 # 10.29.0
 - [feature] Added community support for watchOS. (#13215)

FirebaseVertexAI/Sample/ChatSample/Screens/ConversationScreen.swift

@@ -104,7 +104,9 @@ struct ConversationScreen: View {
   }
 
   private func newChat() {
-    viewModel.startNewChat()
+    Task {
+      await viewModel.startNewChat()
+    }
   }
 }
 

FirebaseVertexAI/Sample/ChatSample/ViewModels/ConversationViewModel.swift

@@ -21,39 +21,44 @@ class ConversationViewModel: ObservableObject {
   /// This array holds both the user's and the system's chat messages
   @Published var messages = [ChatMessage]()
 
-  /// Indicates we're waiting for the model to finish
-  @Published var busy = false
+  /// Indicates we're waiting for the model to finish or the UI is loading
+  @Published var busy = true
 
   @Published var error: Error?
   var hasError: Bool {
     return error != nil
   }
 
   private var model: GenerativeModel
-  private var chat: Chat
+  private var chat: Chat? = nil
   private var stopGenerating = false
 
   private var chatTask: Task<Void, Never>?
 
   init() {
     model = VertexAI.vertexAI().generativeModel(modelName: "gemini-1.5-flash")
-    chat = model.startChat()
+    Task {
+      await startNewChat()
+    }
   }
 
   func sendMessage(_ text: String, streaming: Bool = true) async {
-    error = nil
+    stop()
     if streaming {
       await internalSendMessageStreaming(text)
     } else {
       await internalSendMessage(text)
     }
   }
 
-  func startNewChat() {
+  func startNewChat() async {
+    busy = true
+    defer {
+      busy = false
+    }
     stop()
-    error = nil
-    chat = model.startChat()
     messages.removeAll()
+    chat = await model.startChat()
   }
 
   func stop() {
@@ -62,8 +67,6 @@ class ConversationViewModel: ObservableObject {
   }
 
   private func internalSendMessageStreaming(_ text: String) async {
-    chatTask?.cancel()
-
     chatTask = Task {
       busy = true
       defer {
@@ -79,7 +82,10 @@ class ConversationViewModel: ObservableObject {
       messages.append(systemMessage)
 
       do {
-        let responseStream = chat.sendMessageStream(text)
+        guard let chat else {
+          throw ChatError.notInitialized
+        }
+        let responseStream = await chat.sendMessageStream(text)
         for try await chunk in responseStream {
           messages[messages.count - 1].pending = false
           if let text = chunk.text {
@@ -95,8 +101,6 @@ class ConversationViewModel: ObservableObject {
   }
 
   private func internalSendMessage(_ text: String) async {
-    chatTask?.cancel()
-
     chatTask = Task {
       busy = true
       defer {
@@ -112,10 +116,12 @@ class ConversationViewModel: ObservableObject {
       messages.append(systemMessage)
 
       do {
-        var response: GenerateContentResponse?
-        response = try await chat.sendMessage(text)
+        guard let chat = chat else {
+          throw ChatError.notInitialized
+        }
+        let response = try await chat.sendMessage(text)
 
-        if let responseText = response?.text {
+        if let responseText = response.text {
           // replace pending message with backend response
           messages[messages.count - 1].message = responseText
           messages[messages.count - 1].pending = false
@@ -127,4 +133,8 @@ class ConversationViewModel: ObservableObject {
       }
     }
   }
+
+  enum ChatError: Error {
+    case notInitialized
+  }
 }

FirebaseVertexAI/Sample/FunctionCallingSample/Screens/FunctionCallingScreen.swift

@@ -106,7 +106,9 @@ struct FunctionCallingScreen: View {
   }
 
   private func newChat() {
-    viewModel.startNewChat()
+    Task {
+      await viewModel.startNewChat()
+    }
   }
 }
 

FirebaseVertexAI/Sample/FunctionCallingSample/ViewModels/FunctionCallingViewModel.swift

@@ -33,7 +33,7 @@ class FunctionCallingViewModel: ObservableObject {
   private var functionCalls = [FunctionCall]()
 
   private var model: GenerativeModel
-  private var chat: Chat
+  private var chat: Chat? = nil
 
   private var chatTask: Task<Void, Never>?
 
@@ -62,13 +62,13 @@ class FunctionCallingViewModel: ObservableObject {
         ),
       ])]
     )
-    chat = model.startChat()
+    Task {
+      await startNewChat()
+    }
   }
 
   func sendMessage(_ text: String, streaming: Bool = true) async {
-    error = nil
-    chatTask?.cancel()
-
+    stop()
     chatTask = Task {
       busy = true
       defer {
@@ -100,11 +100,14 @@ class FunctionCallingViewModel: ObservableObject {
     }
   }
 
-  func startNewChat() {
+  func startNewChat() async {
+    busy = true
+    defer {
+      busy = false
+    }
     stop()
-    error = nil
-    chat = model.startChat()
     messages.removeAll()
+    chat = await model.startChat()
   }
 
   func stop() {
@@ -114,14 +117,17 @@ class FunctionCallingViewModel: ObservableObject {
 
   private func internalSendMessageStreaming(_ text: String) async throws {
     let functionResponses = try await processFunctionCalls()
+    guard let chat else {
+      throw ChatError.notInitialized
+    }
     let responseStream: AsyncThrowingStream<GenerateContentResponse, Error>
     if functionResponses.isEmpty {
-      responseStream = chat.sendMessageStream(text)
+      responseStream = await chat.sendMessageStream(text)
     } else {
       for functionResponse in functionResponses {
         messages.insert(functionResponse.chatMessage(), at: messages.count - 1)
       }
-      responseStream = chat.sendMessageStream(functionResponses.modelContent())
+      responseStream = await chat.sendMessageStream(functionResponses.modelContent())
     }
     for try await chunk in responseStream {
       processResponseContent(content: chunk)
@@ -130,6 +136,9 @@ class FunctionCallingViewModel: ObservableObject {
 
   private func internalSendMessage(_ text: String) async throws {
     let functionResponses = try await processFunctionCalls()
+    guard let chat else {
+      throw ChatError.notInitialized
+    }
     let response: GenerateContentResponse
     if functionResponses.isEmpty {
       response = try await chat.sendMessage(text)
@@ -181,6 +190,10 @@ class FunctionCallingViewModel: ObservableObject {
     return functionResponses
   }
 
+  enum ChatError: Error {
+    case notInitialized
+  }
+
   // MARK: - Callable Functions
 
   func getExchangeRate(args: JSONObject) -> JSONObject {

FirebaseVertexAI/Sample/GenerativeAIMultimodalSample/ViewModels/PhotoReasoningViewModel.swift

@@ -84,7 +84,7 @@ class PhotoReasoningViewModel: ObservableObject {
         }
       }
 
-      let outputContentStream = model.generateContentStream(prompt, images)
+      let outputContentStream = await model.generateContentStream(prompt, images)
 
       // stream response
       for try await outputContent in outputContentStream {

FirebaseVertexAI/Sample/GenerativeAITextSample/ViewModels/SummarizeViewModel.swift

@@ -50,7 +50,7 @@ class SummarizeViewModel: ObservableObject {
 
       let prompt = "Summarize the following text for me: \(inputText)"
 
-      let outputContentStream = model.generateContentStream(prompt)
+      let outputContentStream = await model.generateContentStream(prompt)
 
       // stream response
       for try await outputContent in outputContentStream {

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
 
@@ -217,33 +217,31 @@ public final class GenerativeModel {
                                                         isStreaming: true,
                                                         options: requestOptions)
 
-    var responseIterator = generativeAIService.loadRequestStream(request: generateContentRequest)
-      .makeAsyncIterator()
+    let responseStream = generativeAIService.loadRequestStream(request: generateContentRequest)
+
     return AsyncThrowingStream {
-      let response: GenerateContentResponse?
       do {
-        response = try await responseIterator.next()
-      } catch {
-        throw GenerativeModel.generateContentError(from: error)
-      }
+        for try await response in responseStream {
+          // Check the prompt feedback to see if the prompt was blocked.
+          if response.promptFeedback?.blockReason != nil {
+            throw GenerateContentError.promptBlocked(response: response)
+          }
 
-      // The responseIterator will return `nil` when it's done.
-      guard let response = response else {
+          // 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 sending `nil`.
         return nil
-      }
-
-      // Check the prompt feedback to see if the prompt was blocked.
-      if response.promptFeedback?.blockReason != nil {
-        throw GenerateContentError.promptBlocked(response: 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
+      } catch {
+        throw GenerativeModel.generateContentError(from: error)
       }
     }
   }

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)
   }
 }