[FirebaseAI] Add support for Grounding with Google Search (#15014)

Author: dlarocque

FirebaseAI/CHANGELOG.md

@@ -1,4 +1,5 @@
 # Unreleased
+- [added] Added support for Grounding with Google Search. (#15014)
 - [removed] Removed `CountTokensResponse.totalBillableCharacters` which was
   deprecated in 11.15.0. Use `totalTokens` instead. (#15056)
 

FirebaseAI/Sources/GenerateContentResponse.swift

@@ -136,13 +136,16 @@ public struct Candidate: Sendable {
   /// Cited works in the model's response content, if it exists.
   public let citationMetadata: CitationMetadata?
 
+  public let groundingMetadata: GroundingMetadata?
+
   /// Initializer for SwiftUI previews or tests.
   public init(content: ModelContent, safetyRatings: [SafetyRating], finishReason: FinishReason?,
-              citationMetadata: CitationMetadata?) {
+              citationMetadata: CitationMetadata?, groundingMetadata: GroundingMetadata? = nil) {
     self.content = content
     self.safetyRatings = safetyRatings
     self.finishReason = finishReason
     self.citationMetadata = citationMetadata
+    self.groundingMetadata = groundingMetadata
   }
 }
 
@@ -299,6 +302,110 @@ public struct PromptFeedback: Sendable {
   }
 }
 
+/// Metadata returned to the client when grounding is enabled.
+///
+/// > Important: If using Grounding with Google Search, you are required to comply with the
+/// "Grounding with Google Search" usage requirements for your chosen API provider:
+/// [Gemini Developer API](https://ai.google.dev/gemini-api/terms#grounding-with-google-search)
+/// or Vertex AI Gemini API (see [Service Terms](https://cloud.google.com/terms/service-terms)
+/// section within the Service Specific Terms).
+@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
+public struct GroundingMetadata: Sendable {
+  /// A list of web search queries that the model performed to gather the grounding information.
+  /// These can be used to allow users to explore the search results themselves.
+  public let webSearchQueries: [String]
+  /// A list of ``GroundingChunk`` structs. Each chunk represents a piece of retrieved content
+  /// (e.g., from a web page) that the model used to ground its response.
+  public let groundingChunks: [GroundingChunk]
+  /// A list of ``GroundingSupport`` structs. Each object details how specific segments of the
+  /// model's response are supported by the `groundingChunks`.
+  public let groundingSupports: [GroundingSupport]
+  /// Google Search entry point for web searches.
+  /// This contains an HTML/CSS snippet that **must** be embedded in an app to display a Google
+  /// Search entry point for follow-up web searches related to the model's "Grounded Response".
+  public let searchEntryPoint: SearchEntryPoint?
+
+  /// A struct representing the Google Search entry point.
+  @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
+  public struct SearchEntryPoint: Sendable {
+    /// An HTML/CSS snippet that can be embedded in your app.
+    ///
+    /// To ensure proper rendering, it's recommended to display this content within a `WKWebView`.
+    public let renderedContent: String
+  }
+
+  /// Represents a chunk of retrieved data that supports a claim in the model's response. This is
+  /// part of the grounding information provided when grounding is enabled.
+  @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
+  public struct GroundingChunk: Sendable {
+    /// Contains details if the grounding chunk is from a web source.
+    public let web: WebGroundingChunk?
+  }
+
+  /// A grounding chunk sourced from the web.
+  @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
+  public struct WebGroundingChunk: Sendable {
+    /// The URI of the retrieved web page.
+    public let uri: String?
+    /// The title of the retrieved web page.
+    public let title: String?
+    /// The domain of the original URI from which the content was retrieved.
+    ///
+    /// This field is only populated when using the Vertex AI Gemini API.
+    public let domain: String?
+  }
+
+  /// Provides information about how a specific segment of the model's response is supported by the
+  /// retrieved grounding chunks.
+  @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
+  public struct GroundingSupport: Sendable {
+    /// Specifies the segment of the model's response content that this grounding support pertains
+    /// to.
+    public let segment: Segment
+
+    /// A list of indices that refer to specific ``GroundingChunk`` structs within the
+    /// ``GroundingMetadata/groundingChunks`` array. These referenced chunks are the sources that
+    /// support the claim made in the associated `segment` of the response. For example, an array
+    /// `[1, 3, 4]`
+    /// means that `groundingChunks[1]`, `groundingChunks[3]`, `groundingChunks[4]` are the
+    /// retrieved content supporting this part of the response.
+    public let groundingChunkIndices: [Int]
+
+    struct Internal {
+      let segment: Segment?
+      let groundingChunkIndices: [Int]
+
+      func toPublic() -> GroundingSupport? {
+        if segment == nil {
+          return nil
+        }
+        return GroundingSupport(
+          segment: segment!,
+          groundingChunkIndices: groundingChunkIndices
+        )
+      }
+    }
+  }
+}
+
+/// Represents a specific segment within a ``ModelContent`` struct, often used to pinpoint the
+/// exact location of text or data that grounding information refers to.
+@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
+public struct Segment: Sendable {
+  /// The zero-based index of the ``Part`` object within the `parts` array of its parent
+  /// ``ModelContent`` object. This identifies which part of the content the segment belongs to.
+  public let partIndex: Int
+  /// The zero-based start index of the segment within the specified ``Part``, measured in UTF-8
+  /// bytes. This offset is inclusive, starting from 0 at the beginning of the part's content.
+  public let startIndex: Int
+  /// The zero-based end index of the segment within the specified ``Part``, measured in UTF-8
+  /// bytes. This offset is exclusive, meaning the character at this index is not included in the
+  /// segment.
+  public let endIndex: Int
+  /// The text corresponding to the segment from the response.
+  public let text: String
+}
+
 // MARK: - Codable Conformances
 
 @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
@@ -369,6 +476,7 @@ extension Candidate: Decodable {
     case safetyRatings
     case finishReason
     case citationMetadata
+    case groundingMetadata
   }
 
   /// Initializes a response from a decoder. Used for decoding server responses; not for public
@@ -414,6 +522,11 @@ extension Candidate: Decodable {
       CitationMetadata.self,
       forKey: .citationMetadata
     )
+
+    groundingMetadata = try container.decodeIfPresent(
+      GroundingMetadata.self,
+      forKey: .groundingMetadata
+    )
   }
 }
 
@@ -513,3 +626,74 @@ extension PromptFeedback: Decodable {
     }
   }
 }
+
+@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
+extension GroundingMetadata: Decodable {
+  enum CodingKeys: String, CodingKey {
+    case webSearchQueries
+    case groundingChunks
+    case groundingSupports
+    case searchEntryPoint
+  }
+
+  public init(from decoder: Decoder) throws {
+    let container = try decoder.container(keyedBy: CodingKeys.self)
+    webSearchQueries = try container.decodeIfPresent([String].self, forKey: .webSearchQueries) ?? []
+    groundingChunks = try container.decodeIfPresent(
+      [GroundingChunk].self,
+      forKey: .groundingChunks
+    ) ?? []
+    groundingSupports = try container.decodeIfPresent(
+      [GroundingSupport.Internal].self,
+      forKey: .groundingSupports
+    )?.compactMap { $0.toPublic() } ?? []
+    searchEntryPoint = try container.decodeIfPresent(
+      SearchEntryPoint.self,
+      forKey: .searchEntryPoint
+    )
+  }
+}
+
+@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
+extension GroundingMetadata.SearchEntryPoint: Decodable {}
+
+@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
+extension GroundingMetadata.GroundingChunk: Decodable {}
+
+@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
+extension GroundingMetadata.WebGroundingChunk: Decodable {}
+
+@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
+extension GroundingMetadata.GroundingSupport.Internal: Decodable {
+  enum CodingKeys: String, CodingKey {
+    case segment
+    case groundingChunkIndices
+  }
+
+  public init(from decoder: Decoder) throws {
+    let container = try decoder.container(keyedBy: CodingKeys.self)
+    segment = try container.decodeIfPresent(Segment.self, forKey: .segment)
+    groundingChunkIndices = try container.decodeIfPresent(
+      [Int].self,
+      forKey: .groundingChunkIndices
+    ) ?? []
+  }
+}
+
+@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
+extension Segment: Decodable {
+  enum CodingKeys: String, CodingKey {
+    case partIndex
+    case startIndex
+    case endIndex
+    case text
+  }
+
+  public init(from decoder: Decoder) throws {
+    let container = try decoder.container(keyedBy: CodingKeys.self)
+    partIndex = try container.decodeIfPresent(Int.self, forKey: .partIndex) ?? 0
+    startIndex = try container.decodeIfPresent(Int.self, forKey: .startIndex) ?? 0
+    endIndex = try container.decodeIfPresent(Int.self, forKey: .endIndex) ?? 0
+    text = try container.decodeIfPresent(String.self, forKey: .text) ?? ""
+  }
+}

FirebaseAI/Sources/FunctionCalling.swift renamed to FirebaseAI/Sources/Tool.swift

@@ -50,6 +50,19 @@ public struct FunctionDeclaration: Sendable {
   }
 }
 
+/// A tool that allows the generative model to connect to Google Search to access and incorporate
+/// up-to-date information from the web into its responses.
+///
+/// > Important: When using this feature, you are required to comply with the
+/// "Grounding with Google Search" usage requirements for your chosen API provider:
+/// [Gemini Developer API](https://ai.google.dev/gemini-api/terms#grounding-with-google-search)
+/// or Vertex AI Gemini API (see [Service Terms](https://cloud.google.com/terms/service-terms)
+/// section within the Service Specific Terms).
+@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
+public struct GoogleSearch: Sendable {
+  public init() {}
+}
+
 /// A helper tool that the model may use when generating responses.
 ///
 /// A `Tool` is a piece of code that enables the system to interact with external systems to perform
@@ -58,9 +71,17 @@ public struct FunctionDeclaration: Sendable {
 public struct Tool: Sendable {
   /// A list of `FunctionDeclarations` available to the model.
   let functionDeclarations: [FunctionDeclaration]?
+  /// Specifies the Google Search configuration.
+  let googleSearch: GoogleSearch?
 
   init(functionDeclarations: [FunctionDeclaration]?) {
     self.functionDeclarations = functionDeclarations
+    googleSearch = nil
+  }
+
+  init(googleSearch: GoogleSearch) {
+    self.googleSearch = googleSearch
+    functionDeclarations = nil
   }
 
   /// Creates a tool that allows the model to perform function calling.
@@ -85,6 +106,26 @@ public struct Tool: Sendable {
   public static func functionDeclarations(_ functionDeclarations: [FunctionDeclaration]) -> Tool {
     return self.init(functionDeclarations: functionDeclarations)
   }
+
+  /// Creates a tool that allows the model to use Grounding with Google Search.
+  ///
+  /// Grounding with Google Search can be used to allow the model to connect to Google Search to
+  /// access and incorporate up-to-date information from the web into it's responses.
+  ///
+  /// > Important: When using this feature, you are required to comply with the
+  /// "Grounding with Google Search" usage requirements for your chosen API provider:
+  /// [Gemini Developer API](https://ai.google.dev/gemini-api/terms#grounding-with-google-search)
+  /// or Vertex AI Gemini API (see [Service Terms](https://cloud.google.com/terms/service-terms)
+  /// section within the Service Specific Terms).
+  ///
+  /// - Parameters:
+  ///   - googleSearch: An empty ``GoogleSearch`` object. The presence of this object in the list
+  ///     of tools enables the model to use Google Search.
+  ///
+  /// - Returns: A `Tool` configured for Google Search.
+  public static func googleSearch(_ googleSearch: GoogleSearch = GoogleSearch()) -> Tool {
+    return self.init(googleSearch: googleSearch)
+  }
 }
 
 /// Configuration for specifying function calling behavior.
@@ -170,5 +211,8 @@ extension FunctionCallingConfig: Encodable {}
 @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
 extension FunctionCallingConfig.Mode: Encodable {}
 
+@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
+extension GoogleSearch: Encodable {}
+
 @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
 extension ToolConfig: Encodable {}

FirebaseAI/Tests/TestApp/Tests/Integration/GenerateContentIntegrationTests.swift

@@ -247,6 +247,45 @@ struct GenerateContentIntegrationTests {
     #endif // canImport(UIKit)
   }
 
+  @Test(
+    "generateContent with Google Search returns grounding metadata",
+    arguments: InstanceConfig.allConfigs
+  )
+  func generateContent_withGoogleSearch_succeeds(_ config: InstanceConfig) async throws {
+    let model = FirebaseAI.componentInstance(config).generativeModel(
+      modelName: ModelNames.gemini2Flash,
+      tools: [.googleSearch()]
+    )
+    let prompt = "What is the weather in Toronto today?"
+
+    let response = try await model.generateContent(prompt)
+
+    let candidate = try #require(response.candidates.first)
+    let groundingMetadata = try #require(candidate.groundingMetadata)
+    let searchEntrypoint = try #require(groundingMetadata.searchEntryPoint)
+
+    #expect(!groundingMetadata.webSearchQueries.isEmpty)
+    #expect(!searchEntrypoint.renderedContent.isEmpty)
+    #expect(!groundingMetadata.groundingChunks.isEmpty)
+    #expect(!groundingMetadata.groundingSupports.isEmpty)
+
+    for chunk in groundingMetadata.groundingChunks {
+      #expect(chunk.web != nil)
+    }
+
+    for support in groundingMetadata.groundingSupports {
+      let segment = support.segment
+      #expect(segment.endIndex > segment.startIndex)
+      #expect(!segment.text.isEmpty)
+      #expect(!support.groundingChunkIndices.isEmpty)
+
+      // Ensure indices point to valid chunks
+      for index in support.groundingChunkIndices {
+        #expect(index < groundingMetadata.groundingChunks.count)
+      }
+    }
+  }
+
   // MARK: Streaming Tests
 
   @Test(arguments: [