[FirebaseAI ] Integrate conversationkit

firebaseai/ChatExample/Models/ChatMessage.swift

@@ -14,35 +14,65 @@
 
 import FirebaseAI
 import Foundation
+import ConversationKit
 
-enum Participant {
-  case system
-  case user
-}
+public struct ChatMessage: Message {
+  public let id: UUID = .init()
+  public var content: String?
+  public let imageURL: String?
+  public let participant: Participant
+  public let error: (any Error)?
+  public var pending = false
+  public var groundingMetadata: GroundingMetadata?
 
-struct ChatMessage: Identifiable, Equatable {
-  let id = UUID().uuidString
-  var message: String
-  var groundingMetadata: GroundingMetadata?
-  let participant: Participant
-  var pending = false
+  public init(content: String? = nil, imageURL: String? = nil, participant: Participant,
+              error: (any Error)? = nil, pending: Bool = false) {
+    self.content = content
+    self.imageURL = imageURL
+    self.participant = participant
+    self.error = error
+    self.pending = pending
+  }
 
-  static func pending(participant: Participant) -> ChatMessage {
-    Self(message: "", participant: participant, pending: true)
+  // Protocol-required initializer
+  public init(content: String?, imageURL: String?, participant: Participant) {
+    self.content = content
+    self.imageURL = imageURL
+    self.participant = participant
+    error = nil
   }
+}
 
-  // TODO(andrewheard): Add Equatable conformance to GroundingMetadata and remove this
-  static func == (lhs: ChatMessage, rhs: ChatMessage) -> Bool {
-    lhs.id == rhs.id && lhs.message == rhs.message && lhs.participant == rhs.participant && lhs
-      .pending == rhs.pending
+extension ChatMessage {
+  public static func pending(participant: Participant) -> ChatMessage {
+    Self(content: "", participant: participant, pending: true)
   }
 }
 
+// Implement Equatable and Hashable for ChatMessage (ignore error)
 extension ChatMessage {
+  public static func == (lhs: ChatMessage, rhs: ChatMessage) -> Bool {
+    lhs.id == rhs.id &&
+      lhs.content == rhs.content &&
+      lhs.imageURL == rhs.imageURL &&
+      lhs.participant == rhs.participant
+    // intentionally ignore `error`
+  }
+
+  public func hash(into hasher: inout Hasher) {
+    hasher.combine(id)
+    hasher.combine(content)
+    hasher.combine(imageURL)
+    hasher.combine(participant)
+    // intentionally ignore `error`
+  }
+}
+
+public extension ChatMessage {
   static var samples: [ChatMessage] = [
-    .init(message: "Hello. What can I do for you today?", participant: .system),
-    .init(message: "Show me a simple loop in Swift.", participant: .user),
-    .init(message: """
+    .init(content: "Hello. What can I do for you today?", participant: .other),
+    .init(content: "Show me a simple loop in Swift.", participant: .user),
+    .init(content: """
     Sure, here is a simple loop in Swift:
 
     # Example 1
@@ -65,23 +95,23 @@ extension ChatMessage {
     ```
 
     This loop calculates the sum of the numbers from 1 to 100. The variable sum is initialized to 0, and then the for loop iterates over the range of numbers from 1 to 100. The variable i is assigned each number in the range, and the value of i is added to the sum variable. After the loop has finished executing, the value of sum is printed to the console.
-    """, participant: .system),
+    """, participant: .other),
   ]
 
   static var sample = samples[0]
 }
 
-extension ChatMessage {
+public extension ChatMessage {
   static func from(_ modelContent: ModelContent) -> ChatMessage? {
     // TODO: add non-text parts to message when multi-model support is added
     let text = modelContent.parts.compactMap { ($0 as? TextPart)?.text }.joined()
     guard !text.isEmpty else {
       return nil
     }
 
-    let participant: Participant = (modelContent.role == "user") ? .user : .system
+    let participant: Participant = (modelContent.role == "user") ? .user : .other
 
-    return ChatMessage(message: text, participant: participant)
+    return ChatMessage(content: text, participant: participant)
   }
 
   static func from(_ modelContents: [ModelContent]) -> [ChatMessage] {

firebaseai/ChatExample/Screens/ChatScreen.swift

@@ -0,0 +1,69 @@
+// Copyright 2023 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+import FirebaseAI
+import SwiftUI
+import ConversationKit
+
+struct ChatScreen: View {
+  let firebaseService: FirebaseAI
+  @StateObject var viewModel: ChatViewModel
+
+  init(firebaseService: FirebaseAI, sample: Sample? = nil) {
+    self.firebaseService = firebaseService
+    _viewModel =
+      StateObject(wrappedValue: ChatViewModel(firebaseService: firebaseService,
+                                              sample: sample))
+  }
+
+  var body: some View {
+    NavigationStack {
+      ConversationView(messages: $viewModel.messages,
+                       userPrompt: viewModel.initialPrompt) { message in
+        MessageView(message: message)
+      }
+      .disableAttachments()
+      .onSendMessage { message in
+        Task {
+          await viewModel.sendMessage(message.content ?? "", streaming: true)
+        }
+      }
+      .onError { error in
+        viewModel.presentErrorDetails = true
+      }
+      .sheet(isPresented: $viewModel.presentErrorDetails) {
+        if let error = viewModel.error {
+          ErrorDetailsView(error: error)
+        }
+      }
+      .toolbar {
+        ToolbarItem(placement: .primaryAction) {
+          Button(action: newChat) {
+            Image(systemName: "square.and.pencil")
+          }
+        }
+      }
+      .navigationTitle(viewModel.title)
+      .navigationBarTitleDisplayMode(.inline)
+    }
+  }
+
+  private func newChat() {
+    viewModel.startNewChat()
+  }
+}
+
+#Preview {
+  ChatScreen(firebaseService: FirebaseAI.firebaseAI())
+}