Revert "Make GenerativeModel and Chat into Swift actors (#13545)" (#13703)

Author: andrewheard

FirebaseVertexAI/CHANGELOG.md

@@ -1,5 +1,9 @@
-# Unreleased
+# 11.3.0
 - [added] Added `Decodable` conformance for `FunctionResponse`. (#13606)
+- [changed] **Breaking Change**: Reverted refactor of `GenerativeModel` and
+  `Chat` as Swift actors (#13545) introduced in 11.2; The methods
+  `generateContentStream`, `startChat` and `sendMessageStream` no longer need to
+  be called with `await`. (#13703)
 
 # 11.2.0
 - [fixed] Resolved a decoding error for citations without a `uri` and added

FirebaseVertexAI/Sample/ChatSample/Screens/ConversationScreen.swift

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

FirebaseVertexAI/Sample/ChatSample/ViewModels/ConversationViewModel.swift

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

FirebaseVertexAI/Sample/FunctionCallingSample/Screens/FunctionCallingScreen.swift

@@ -106,9 +106,7 @@ struct FunctionCallingScreen: View {
   }
 
   private func newChat() {
-    Task {
-      await viewModel.startNewChat()
-    }
+    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? = nil
+  private var chat: Chat
 
   private var chatTask: Task<Void, Never>?
 
@@ -62,13 +62,13 @@ class FunctionCallingViewModel: ObservableObject {
         ),
       ])]
     )
-    Task {
-      await startNewChat()
-    }
+    chat = model.startChat()
   }
 
   func sendMessage(_ text: String, streaming: Bool = true) async {
-    stop()
+    error = nil
+    chatTask?.cancel()
+
     chatTask = Task {
       busy = true
       defer {
@@ -100,14 +100,11 @@ class FunctionCallingViewModel: ObservableObject {
     }
   }
 
-  func startNewChat() async {
-    busy = true
-    defer {
-      busy = false
-    }
+  func startNewChat() {
     stop()
+    error = nil
+    chat = model.startChat()
     messages.removeAll()
-    chat = await model.startChat()
   }
 
   func stop() {
@@ -117,17 +114,14 @@ 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 = try await chat.sendMessageStream(text)
+      responseStream = try chat.sendMessageStream(text)
     } else {
       for functionResponse in functionResponses {
         messages.insert(functionResponse.chatMessage(), at: messages.count - 1)
       }
-      responseStream = try await chat.sendMessageStream(functionResponses.modelContent())
+      responseStream = try chat.sendMessageStream(functionResponses.modelContent())
     }
     for try await chunk in responseStream {
       processResponseContent(content: chunk)
@@ -136,9 +130,6 @@ 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)
@@ -190,10 +181,6 @@ 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 = try await model.generateContentStream(prompt, images)
+      let outputContentStream = try 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 = try await model.generateContentStream(prompt)
+      let outputContentStream = try 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 actor Chat {
+public class Chat {
   private let model: GenerativeModel
 
   /// Initializes a new chat representing a 1:1 conversation between model and user.
@@ -116,7 +116,7 @@ public actor Chat {
 
         // Send the history alongside the new message as context.
         let request = history + newContent
-        let stream = try await model.generateContentStream(request)
+        let stream = try 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 actor GenerativeModel {
+public final class GenerativeModel {
   /// The resource name of the model in the backend; has the format "models/model-name".
   let modelResourceName: String
 
@@ -212,31 +212,33 @@ public final actor GenerativeModel {
                                                         isStreaming: true,
                                                         options: requestOptions)
 
-    let responseStream = generativeAIService.loadRequestStream(request: generateContentRequest)
-
+    var responseIterator = generativeAIService.loadRequestStream(request: generateContentRequest)
+      .makeAsyncIterator()
     return AsyncThrowingStream {
+      let response: GenerateContentResponse?
       do {
-        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)
-          }
+        response = try await responseIterator.next()
+      } catch {
+        throw GenerativeModel.generateContentError(from: error)
+      }
 
-          // 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
-          }
-        }
+      // 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
-      } catch {
-        throw GenerativeModel.generateContentError(from: error)
+      }
+
+      // 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
       }
     }
   }

FirebaseVertexAI/Tests/Unit/ChatTests.swift

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