Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view

This file was deleted.

This file was deleted.

19 changes: 13 additions & 6 deletions firebaseai/ChatExample/Models/ChatMessage.swift
Original file line number Diff line number Diff line change
Expand Up @@ -15,27 +15,34 @@
import FirebaseAI
import Foundation
import ConversationKit
import UIKit

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?
public var attachments: [MultimodalAttachment] = []
public var image: UIImage?
// required by the Message protocol, but not used in this app
public var imageURL: String?

public init(content: String? = nil, imageURL: String? = nil, participant: Participant,
error: (any Error)? = nil, pending: Bool = false) {
error: (any Error)? = nil, pending: Bool = false,
attachments: [MultimodalAttachment] = [], image: UIImage? = nil) {
self.content = content
self.imageURL = imageURL
self.participant = participant
self.error = error
self.pending = pending
self.attachments = attachments
self.image = image
}

// Protocol-required initializer
public init(content: String?, imageURL: String?, participant: Participant) {
public init(content: String?, imageURL: String? = nil, participant: Participant) {
self.content = content
self.imageURL = imageURL
self.participant = participant
Expand All @@ -54,16 +61,16 @@ 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
lhs.participant == rhs.participant &&
lhs.image == rhs.image
// intentionally ignore `error`
}

public func hash(into hasher: inout Hasher) {
hasher.combine(id)
hasher.combine(content)
hasher.combine(imageURL)
hasher.combine(participant)
hasher.combine(image)
// intentionally ignore `error`
}
}
Expand Down
10 changes: 5 additions & 5 deletions firebaseai/ChatExample/Screens/ChatScreen.swift
Original file line number Diff line number Diff line change
Expand Up @@ -17,13 +17,13 @@ import SwiftUI
import ConversationKit

struct ChatScreen: View {
let firebaseService: FirebaseAI
let backendType: BackendOption
@StateObject var viewModel: ChatViewModel

init(firebaseService: FirebaseAI, sample: Sample? = nil) {
self.firebaseService = firebaseService
init(backendType: BackendOption, sample: Sample? = nil) {
self.backendType = backendType
_viewModel =
StateObject(wrappedValue: ChatViewModel(firebaseService: firebaseService,
StateObject(wrappedValue: ChatViewModel(backendType: backendType,
sample: sample))
}

Expand Down Expand Up @@ -65,5 +65,5 @@ struct ChatScreen: View {
}

#Preview {
ChatScreen(firebaseService: FirebaseAI.firebaseAI())
ChatScreen(backendType: .googleAI)
}
38 changes: 24 additions & 14 deletions firebaseai/ChatExample/ViewModels/ChatViewModel.swift
Original file line number Diff line number Diff line change
Expand Up @@ -36,24 +36,31 @@ class ChatViewModel: ObservableObject {

private var model: GenerativeModel
private var chat: Chat
private var stopGenerating = false

private var chatTask: Task<Void, Never>?

private var sample: Sample?
private var backendType: BackendOption

init(firebaseService: FirebaseAI, sample: Sample? = nil) {
init(backendType: BackendOption, sample: Sample? = nil) {
self.sample = sample
self.backendType = backendType

let firebaseService: FirebaseAI
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think a ternary might be more compact:

    let firebaseService = backendType == .googleAI
      ? FirebaseAI.firebaseAI(backend: .googleAI())
      : FirebaseAI.firebaseAI(backend: .vertexAI())

switch backendType {
case .googleAI:
firebaseService = FirebaseAI.firebaseAI(backend: .googleAI())
case .vertexAI:
firebaseService = FirebaseAI.firebaseAI(backend: .vertexAI())
}

// create a generative model with sample data
model = firebaseService.generativeModel(
modelName: "gemini-2.0-flash-001",
tools: sample?.tools,
modelName: sample?.modelName ?? "gemini-2.5-flash",
generationConfig: sample?.generationConfig,
systemInstruction: sample?.systemInstruction
)

if let chatHistory = sample?.chatHistory, !chatHistory.isEmpty {
// Initialize with sample chat history if it's available
messages = ChatMessage.from(chatHistory)
chat = model.startChat(history: chatHistory)
} else {
Expand Down Expand Up @@ -112,13 +119,14 @@ class ChatViewModel: ObservableObject {
.content = (messages[messages.count - 1].content ?? "") + text
}

if let candidate = chunk.candidates.first {
if let groundingMetadata = candidate.groundingMetadata {
self.messages[self.messages.count - 1].groundingMetadata = groundingMetadata
if let inlineDataPart = chunk.inlineDataParts.first {
if let uiImage = UIImage(data: inlineDataPart.data) {
messages[messages.count - 1].image = uiImage
} else {
print("Failed to convert inline data to UIImage")
}
}
Comment on lines +122 to 128

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

This block of code for handling an inlineDataPart is duplicated in internalSendMessage on lines 159-165. To improve maintainability and reduce redundancy, consider extracting this logic into a private helper method. This would make the code cleaner and easier to manage.

}

} catch {
self.error = error
print(error.localizedDescription)
Expand Down Expand Up @@ -156,11 +164,13 @@ class ChatViewModel: ObservableObject {
// replace pending message with backend response
messages[messages.count - 1].content = responseText
messages[messages.count - 1].pending = false
}

if let candidate = response?.candidates.first {
if let groundingMetadata = candidate.groundingMetadata {
self.messages[self.messages.count - 1].groundingMetadata = groundingMetadata
}
if let inlineDataPart = response?.inlineDataParts.first {
if let uiImage = UIImage(data: inlineDataPart.data) {
messages[messages.count - 1].image = uiImage
} else {
print("Failed to convert inline data to UIImage")
}
}
} catch {
Expand Down
28 changes: 20 additions & 8 deletions firebaseai/ChatExample/Views/MessageView.swift
Original file line number Diff line number Diff line change
Expand Up @@ -53,16 +53,28 @@ struct MessageContentView: View {
}
.labelStyle(.iconOnly)
}
}
} else {
VStack(alignment: .leading, spacing: 8) {
if message.participant == .user && !message.attachments.isEmpty {
AttachmentPreviewScrollView(attachments: message.attachments)
}

// Grounded Response
else if let groundingMetadata = message.groundingMetadata {
GroundedResponseView(message: message, groundingMetadata: groundingMetadata)
}
if let image = message.image {
Image(uiImage: image)
.resizable()
.aspectRatio(contentMode: .fit)
.frame(maxWidth: 300, maxHeight: 300)
.clipShape(RoundedRectangle(cornerRadius: 8))
}

// Non-grounded response
else {
ResponseTextView(message: message)
// Grounded Response
if let groundingMetadata = message.groundingMetadata {
GroundedResponseView(message: message, groundingMetadata: groundingMetadata)
} else {
// Non-grounded response
ResponseTextView(message: message)
}
}
}
}
}
Expand Down
Loading