[Firebase AI] Add handling for Google AI-formatted CitationMetadata

FirebaseAI/Sources/GenerateContentResponse.swift

@@ -145,7 +145,7 @@ public struct CitationMetadata: Sendable {
 
 /// A struct describing a source attribution.
 @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
-public struct Citation: Sendable {
+public struct Citation: Sendable, Equatable {
   /// The inclusive beginning of a sequence in a model response that derives from a cited source.
   public let startIndex: Int
 
@@ -165,6 +165,20 @@ public struct Citation: Sendable {
   ///
   /// > Tip: `DateComponents` can be converted to a `Date` using the `date` computed property.
   public let publicationDate: DateComponents?
+
+  init(startIndex: Int,
+       endIndex: Int,
+       uri: String? = nil,
+       title: String? = nil,
+       license: String? = nil,
+       publicationDate: DateComponents? = nil) {
+    self.startIndex = startIndex
+    self.endIndex = endIndex
+    self.uri = uri
+    self.title = title
+    self.license = license
+    self.publicationDate = publicationDate
+  }
 }
 
 /// A value enumerating possible reasons for a model to terminate a content generation request.
@@ -385,7 +399,23 @@ extension Candidate: Decodable {
 }
 
 @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
-extension CitationMetadata: Decodable {}
+extension CitationMetadata: Decodable {
+  enum CodingKeys: CodingKey {
+    case citations // Vertex AI
+    case citationSources // Google AI
+  }
+
+  public init(from decoder: any Decoder) throws {
+    let container = try decoder.container(keyedBy: CodingKeys.self)
+
+    // Decode for Google API if `citationSources` key is present.
+    if container.contains(.citationSources) {
+      citations = try container.decode([Citation].self, forKey: .citationSources)
+    } else { // Fallback to default Vertex AI decoding.
+      citations = try container.decode([Citation].self, forKey: .citations)
+    }
+  }
+}
 
 @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
 extension Citation: Decodable {

FirebaseAI/Tests/Unit/Types/CitationMetadataTests.swift

@@ -0,0 +1,78 @@
+// Copyright 2025 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 XCTest
+
+@testable import FirebaseAI
+
+@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
+final class CitationMetadataTests: XCTestCase {
+  let decoder = JSONDecoder()
+
+  let expectedStartIndex = 100
+  let expectedEndIndex = 200
+  let expectedURI = "https://example.com/citation-1"
+  lazy var citationJSON = """
+  {
+    "startIndex" : \(expectedStartIndex),
+    "endIndex" : \(expectedEndIndex),
+    "uri" : "\(expectedURI)"
+  }
+  """
+  lazy var expectedCitation = Citation(
+    startIndex: expectedStartIndex,
+    endIndex: expectedEndIndex,
+    uri: expectedURI
+  )
+
+  // MARK: - Google AI Format Decoding
+
+  func testDecodeCitationMetadata_googleAIFormat() throws {
+    let json = """
+    {
+      "citationSources": [\(citationJSON)]
+    }
+    """
+    let jsonData = try XCTUnwrap(json.data(using: .utf8))
+
+    let citationMetadata = try decoder.decode(
+      CitationMetadata.self,
+      from: jsonData
+    )
+
+    XCTAssertEqual(citationMetadata.citations.count, 1)
+    let citation = try XCTUnwrap(citationMetadata.citations.first)
+    XCTAssertEqual(citation, expectedCitation)
+  }
+
+  // MARK: - Vertex AI Format Decoding
+
+  func testDecodeCitationMetadata_vertexAIFormat() throws {
+    let json = """
+    {
+      "citations": [\(citationJSON)]
+    }
+    """
+    let jsonData = try XCTUnwrap(json.data(using: .utf8))
+
+    let citationMetadata = try decoder.decode(
+      CitationMetadata.self,
+      from: jsonData
+    )
+
+    XCTAssertEqual(citationMetadata.citations.count, 1)
+    let citation = try XCTUnwrap(citationMetadata.citations.first)
+    XCTAssertEqual(citation, expectedCitation)
+  }
+}

FirebaseAI/Tests/Unit/Types/CitationTests.swift

@@ -0,0 +1,113 @@
+// Copyright 2025 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 XCTest
+
+@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
+final class CitationTests: XCTestCase {
+  let decoder = JSONDecoder()
+
+  // MARK: - Decoding Tests
+
+  func testDecodeCitation_minimalParameters() throws {
+    let expectedEndIndex = 150
+    let json = """
+    {
+      "endIndex" : \(expectedEndIndex)
+    }
+    """
+    let jsonData = try XCTUnwrap(json.data(using: .utf8))
+
+    let citation = try decoder.decode(Citation.self, from: jsonData)
+
+    XCTAssertEqual(citation.startIndex, 0, "Omitted startIndex should be decoded as 0.")
+    XCTAssertEqual(citation.endIndex, expectedEndIndex)
+    XCTAssertNil(citation.uri)
+    XCTAssertNil(citation.title)
+    XCTAssertNil(citation.license)
+    XCTAssertNil(citation.publicationDate)
+  }
+
+  func testDecodeCitation_allParameters() throws {
+    let expectedStartIndex = 100
+    let expectedEndIndex = 200
+    let expectedURI = "https://example.com/citation-1"
+    let expectedTitle = "Example Citation Title"
+    let expectedLicense = "mit"
+    let expectedYear = 2023
+    let expectedMonth = 10
+    let expectedDay = 26
+    let json = """
+    {
+      "startIndex" : \(expectedStartIndex),
+      "endIndex" : \(expectedEndIndex),
+      "uri" : "\(expectedURI)",
+      "title" : "\(expectedTitle)",
+      "license" : "\(expectedLicense)",
+      "publicationDate" : {
+        "year" : \(expectedYear),
+        "month" : \(expectedMonth),
+        "day" : \(expectedDay)
+      }
+    }
+    """
+    let jsonData = try XCTUnwrap(json.data(using: .utf8))
+
+    let citation = try decoder.decode(Citation.self, from: jsonData)
+
+    XCTAssertEqual(citation.startIndex, expectedStartIndex)
+    XCTAssertEqual(citation.endIndex, expectedEndIndex)
+    XCTAssertEqual(citation.uri, expectedURI)
+    XCTAssertEqual(citation.title, expectedTitle)
+    XCTAssertEqual(citation.license, expectedLicense)
+    let publicationDate = try XCTUnwrap(citation.publicationDate)
+    XCTAssertEqual(publicationDate.year, expectedYear)
+    XCTAssertEqual(publicationDate.month, expectedMonth)
+    XCTAssertEqual(publicationDate.day, expectedDay)
+  }
+
+  func testDecodeCitation_emptyStringsForOptionals_setsToNil() throws {
+    let expectedEndIndex = 300
+    let json = """
+    {
+      "endIndex" : \(expectedEndIndex),
+      "uri" : "",
+      "title" : "",
+      "license" : ""
+    }
+    """
+    let jsonData = try XCTUnwrap(json.data(using: .utf8))
+
+    let citation = try decoder.decode(Citation.self, from: jsonData)
+
+    XCTAssertEqual(citation.startIndex, 0, "Omitted startIndex should be decoded as 0.")
+    XCTAssertEqual(citation.endIndex, expectedEndIndex)
+    XCTAssertNil(citation.uri, "Empty URI string should be decoded as nil.")
+    XCTAssertNil(citation.title, "Empty title string should be decoded as nil.")
+    XCTAssertNil(citation.license, "Empty license string should be decoded as nil.")
+    XCTAssertNil(citation.publicationDate)
+  }
+
+  func testDecodeCitation_missingEndIndex_throws() throws {
+    let json = """
+    {
+      "startIndex" : 10
+    }
+    """
+    let jsonData = try XCTUnwrap(json.data(using: .utf8))
+
+    XCTAssertThrowsError(try decoder.decode(Citation.self, from: jsonData))
+  }
+}