Skip to content

Commit 979388a

Browse files
dlarocquedlarocque
authored
[FirebaseAI] Add usage of Grounding with Google Search (#1724)
1 parent 212a9f2 commit 979388a

File tree

8 files changed

+263
-29
lines changed

8 files changed

+263
-29
lines changed

‎firebaseai/ChatExample/Models/ChatMessage.swift

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
// See the License for the specific language governing permissions and
1313
// limitations under the License.
1414

15+
import FirebaseAI
1516
import Foundation
1617

1718
enum Participant {
@@ -22,12 +23,19 @@ enum Participant {
2223
struct ChatMessage: Identifiable, Equatable {
2324
let id = UUID().uuidString
2425
var message: String
26+
var groundingMetadata: GroundingMetadata?
2527
let participant: Participant
2628
var pending = false
2729

2830
static func pending(participant: Participant) -> ChatMessage {
2931
Self(message: "", participant: participant, pending: true)
3032
}
33+
34+
// TODO(andrewheard): Add Equatable conformance to GroundingMetadata and remove this
35+
static func == (lhs: ChatMessage, rhs: ChatMessage) -> Bool {
36+
lhs.id == rhs.id && lhs.message == rhs.message && lhs.participant == rhs.participant && lhs
37+
.pending == rhs.pending
38+
}
3139
}
3240

3341
extension ChatMessage {

‎firebaseai/ChatExample/Screens/ConversationScreen.swift

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -18,15 +18,22 @@ import SwiftUI
1818

1919
struct ConversationScreen: View {
2020
let firebaseService: FirebaseAI
21+
let title: String
2122
@StateObject var viewModel: ConversationViewModel
2223

2324
@State
2425
private var userPrompt = ""
2526

26-
init(firebaseService: FirebaseAI) {
27+
init(firebaseService: FirebaseAI, title: String, searchGroundingEnabled: Bool = false) {
28+
let model = firebaseService.generativeModel(
29+
modelName: "gemini-2.0-flash-001",
30+
tools: searchGroundingEnabled ? [.googleSearch()] : []
31+
)
32+
self.title = title
2733
self.firebaseService = firebaseService
2834
_viewModel =
29-
StateObject(wrappedValue: ConversationViewModel(firebaseService: firebaseService))
35+
StateObject(wrappedValue: ConversationViewModel(firebaseService: firebaseService,
36+
model: model))
3037
}
3138

3239
enum FocusedField: Hashable {
@@ -88,7 +95,7 @@ struct ConversationScreen: View {
8895
}
8996
}
9097
}
91-
.navigationTitle("Chat example")
98+
.navigationTitle(title)
9299
.onAppear {
93100
focusedField = .message
94101
}
@@ -123,7 +130,7 @@ struct ConversationScreen_Previews: PreviewProvider {
123130
.firebaseAI()) // Example service init
124131

125132
var body: some View {
126-
ConversationScreen(firebaseService: FirebaseAI.firebaseAI())
133+
ConversationScreen(firebaseService: FirebaseAI.firebaseAI(), title: "Chat sample")
127134
.onAppear {
128135
viewModel.messages = ChatMessage.samples
129136
}
@@ -132,7 +139,7 @@ struct ConversationScreen_Previews: PreviewProvider {
132139

133140
static var previews: some View {
134141
NavigationStack {
135-
ConversationScreen(firebaseService: FirebaseAI.firebaseAI())
142+
ConversationScreen(firebaseService: FirebaseAI.firebaseAI(), title: "Chat sample")
136143
}
137144
}
138145
}

‎firebaseai/ChatExample/ViewModels/ConversationViewModel.swift

Lines changed: 22 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -35,9 +35,15 @@ class ConversationViewModel: ObservableObject {
3535

3636
private var chatTask: Task<Void, Never>?
3737

38-
init(firebaseService: FirebaseAI) {
39-
model = firebaseService.generativeModel(modelName: "gemini-2.0-flash-001")
40-
chat = model.startChat()
38+
init(firebaseService: FirebaseAI, model: GenerativeModel? = nil) {
39+
if let model {
40+
self.model = model
41+
} else {
42+
self.model = firebaseService.generativeModel(
43+
modelName: "gemini-2.0-flash-001"
44+
)
45+
}
46+
chat = self.model.startChat()
4147
}
4248

4349
func sendMessage(_ text: String, streaming: Bool = true) async {
@@ -85,7 +91,14 @@ class ConversationViewModel: ObservableObject {
8591
if let text = chunk.text {
8692
messages[messages.count - 1].message += text
8793
}
94+
95+
if let candidate = chunk.candidates.first {
96+
if let groundingMetadata = candidate.groundingMetadata {
97+
self.messages[self.messages.count - 1].groundingMetadata = groundingMetadata
98+
}
99+
}
88100
}
101+
89102
} catch {
90103
self.error = error
91104
print(error.localizedDescription)
@@ -119,6 +132,12 @@ class ConversationViewModel: ObservableObject {
119132
// replace pending message with backend response
120133
messages[messages.count - 1].message = responseText
121134
messages[messages.count - 1].pending = false
135+
136+
if let candidate = response?.candidates.first {
137+
if let groundingMetadata = candidate.groundingMetadata {
138+
self.messages[self.messages.count - 1].groundingMetadata = groundingMetadata
139+
}
140+
}
122141
}
123142
} catch {
124143
self.error = error

‎firebaseai/ChatExample/Views/Grounding/GoogleSearchSuggestionView.swift

Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,79 @@
1+
// Copyright 2025 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
import SwiftUI
16+
import WebKit
17+
18+
/// A view that renders Google Search suggestions with links that allow users
19+
/// to view the search results in the device's default browser.
20+
/// This is added to the bottom of chat messages containing results grounded
21+
/// in Google Search.
22+
struct GoogleSearchSuggestionView: UIViewRepresentable {
23+
let htmlString: String
24+
25+
// This Coordinator class will act as the web view's navigation delegate.
26+
class Coordinator: NSObject, WKNavigationDelegate {
27+
func webView(_ webView: WKWebView,
28+
decidePolicyFor navigationAction: WKNavigationAction,
29+
decisionHandler: @escaping (WKNavigationActionPolicy) -> Void) {
30+
// Check if the navigation was triggered by a user clicking a link.
31+
if navigationAction.navigationType == .linkActivated {
32+
if let url = navigationAction.request.url {
33+
// Open the URL in the system's default browser (e.g., Safari).
34+
UIApplication.shared.open(url)
35+
}
36+
// Cancel the navigation inside our small web view.
37+
decisionHandler(.cancel)
38+
return
39+
}
40+
// For all other navigation types (like the initial HTML load), allow it.
41+
decisionHandler(.allow)
42+
}
43+
}
44+
45+
func makeCoordinator() -> Coordinator {
46+
Coordinator()
47+
}
48+
49+
func makeUIView(context: Context) -> WKWebView {
50+
let webView = WKWebView()
51+
webView.isOpaque = false
52+
webView.backgroundColor = .clear
53+
webView.scrollView.backgroundColor = .clear
54+
webView.scrollView.isScrollEnabled = false
55+
// Set the coordinator as the navigation delegate.
56+
webView.navigationDelegate = context.coordinator
57+
return webView
58+
}
59+
60+
func updateUIView(_ uiView: WKWebView, context: Context) {
61+
// The renderedContent is an HTML snippet with CSS.
62+
// For it to render correctly, we wrap it in a basic HTML document structure.
63+
let fullHTML = """
64+
<!DOCTYPE html>
65+
<html>
66+
<head>
67+
<meta name='viewport' content='width=device-width, initial-scale=1.0, user-scalable=no'>
68+
<style>
69+
body { margin: 0; padding: 0; }
70+
</style>
71+
</head>
72+
<body>
73+
\(htmlString)
74+
</body>
75+
</html>
76+
"""
77+
uiView.loadHTMLString(fullHTML, baseURL: nil)
78+
}
79+
}

‎firebaseai/ChatExample/Views/Grounding/GroundedResponseView.swift

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,81 @@
1+
// Copyright 2023 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
import FirebaseAI
16+
import SwiftUI
17+
18+
/// A view that displays a chat message that is grounded in Google Search.
19+
struct GroundedResponseView: View {
20+
var message: ChatMessage
21+
var groundingMetadata: GroundingMetadata
22+
23+
var body: some View {
24+
// We can only display a response grounded in Google Search if the searchEntrypoint is non-nil.
25+
let isCompliant = (groundingMetadata.groundingChunks.isEmpty || groundingMetadata
26+
.searchEntryPoint != nil)
27+
if isCompliant {
28+
HStack(alignment: .top, spacing: 8) {
29+
VStack(alignment: .leading, spacing: 8) {
30+
// Message text
31+
ResponseTextView(message: message)
32+
33+
if !groundingMetadata.groundingChunks.isEmpty {
34+
Divider()
35+
// Source links
36+
ForEach(0 ..< groundingMetadata.groundingChunks.count, id: \.self) { index in
37+
if let webChunk = groundingMetadata.groundingChunks[index].web {
38+
SourceLinkView(
39+
title: webChunk.title ?? "Untitled Source",
40+
uri: webChunk.uri
41+
)
42+
}
43+
}
44+
}
45+
// Search suggestions
46+
if let searchEntryPoint = groundingMetadata.searchEntryPoint {
47+
Divider()
48+
GoogleSearchSuggestionView(htmlString: searchEntryPoint.renderedContent)
49+
.frame(height: 44)
50+
.clipShape(RoundedRectangle(cornerRadius: 22))
51+
}
52+
}
53+
}
54+
.frame(maxWidth: .infinity, alignment: .leading)
55+
}
56+
}
57+
}
58+
59+
/// A view for a single, clickable source link.
60+
struct SourceLinkView: View {
61+
let title: String
62+
let uri: String?
63+
64+
var body: some View {
65+
if let uri, let url = URL(string: uri) {
66+
Link(destination: url) {
67+
HStack(spacing: 4) {
68+
Image(systemName: "link")
69+
.font(.caption)
70+
.foregroundColor(.secondary)
71+
Text(title)
72+
.font(.footnote)
73+
.underline()
74+
.lineLimit(1)
75+
.multilineTextAlignment(.leading)
76+
}
77+
}
78+
.buttonStyle(.plain)
79+
}
80+
}
81+
}

‎firebaseai/ChatExample/Views/MessageView.swift

Lines changed: 34 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414

1515
import MarkdownUI
1616
import SwiftUI
17+
import FirebaseAI
1718

1819
struct RoundedCorner: Shape {
1920
var radius: CGFloat = .infinity
@@ -42,29 +43,43 @@ struct MessageContentView: View {
4243
if message.pending {
4344
BouncingDots()
4445
} else {
45-
Markdown(message.message)
46-
.markdownTextStyle {
47-
FontFamilyVariant(.normal)
48-
FontSize(.em(0.85))
49-
ForegroundColor(message.participant == .system ? Color(UIColor.label) : .white)
50-
}
51-
.markdownBlockStyle(\.codeBlock) { configuration in
52-
configuration.label
53-
.relativeLineSpacing(.em(0.25))
54-
.markdownTextStyle {
55-
FontFamilyVariant(.monospaced)
56-
FontSize(.em(0.85))
57-
ForegroundColor(Color(.label))
58-
}
59-
.padding()
60-
.background(Color(.secondarySystemBackground))
61-
.clipShape(RoundedRectangle(cornerRadius: 8))
62-
.markdownMargin(top: .zero, bottom: .em(0.8))
63-
}
46+
// Grounded Response
47+
if let groundingMetadata = message.groundingMetadata {
48+
GroundedResponseView(message: message, groundingMetadata: groundingMetadata)
49+
} else {
50+
// Non-grounded response
51+
ResponseTextView(message: message)
52+
}
6453
}
6554
}
6655
}
6756

57+
struct ResponseTextView: View {
58+
var message: ChatMessage
59+
60+
var body: some View {
61+
Markdown(message.message)
62+
.markdownTextStyle {
63+
FontFamilyVariant(.normal)
64+
FontSize(.em(0.85))
65+
ForegroundColor(message.participant == .system ? Color(UIColor.label) : .white)
66+
}
67+
.markdownBlockStyle(\.codeBlock) { configuration in
68+
configuration.label
69+
.relativeLineSpacing(.em(0.25))
70+
.markdownTextStyle {
71+
FontFamilyVariant(.monospaced)
72+
FontSize(.em(0.85))
73+
ForegroundColor(Color(.label))
74+
}
75+
.padding()
76+
.background(Color(.secondarySystemBackground))
77+
.clipShape(RoundedRectangle(cornerRadius: 8))
78+
.markdownMargin(top: .zero, bottom: .em(0.8))
79+
}
80+
}
81+
}
82+
6883
struct MessageView: View {
6984
var message: ChatMessage
7085

0 commit comments

Comments
 (0)