[Vertex AI] Add integration tests for modality usage metrics (#14446)

Author: andrewheard

FirebaseVertexAI/Tests/TestApp/Tests/Integration/IntegrationTests.swift

@@ -39,6 +39,8 @@ final class IntegrationTests: XCTestCase {
     SafetySetting(harmCategory: .dangerousContent, threshold: .blockLowAndAbove),
     SafetySetting(harmCategory: .civicIntegrity, threshold: .blockLowAndAbove),
   ]
+  // Candidates and total token counts may differ slightly between runs due to whitespace tokens.
+  let tokenCountAccuracy = 1
 
   var vertex: VertexAI!
   var model: GenerativeModel!
@@ -74,6 +76,18 @@ final class IntegrationTests: XCTestCase {
 
     let text = try XCTUnwrap(response.text).trimmingCharacters(in: .whitespacesAndNewlines)
     XCTAssertEqual(text, "Mountain View")
+    let usageMetadata = try XCTUnwrap(response.usageMetadata)
+    XCTAssertEqual(usageMetadata.promptTokenCount, 21)
+    XCTAssertEqual(usageMetadata.candidatesTokenCount, 3, accuracy: tokenCountAccuracy)
+    XCTAssertEqual(usageMetadata.totalTokenCount, 24, accuracy: tokenCountAccuracy)
+    XCTAssertEqual(usageMetadata.promptTokensDetails.count, 1)
+    let promptTokensDetails = try XCTUnwrap(usageMetadata.promptTokensDetails.first)
+    XCTAssertEqual(promptTokensDetails.modality, .text)
+    XCTAssertEqual(promptTokensDetails.tokenCount, usageMetadata.promptTokenCount)
+    XCTAssertEqual(usageMetadata.candidatesTokensDetails.count, 1)
+    let candidatesTokensDetails = try XCTUnwrap(usageMetadata.candidatesTokensDetails.first)
+    XCTAssertEqual(candidatesTokensDetails.modality, .text)
+    XCTAssertEqual(candidatesTokensDetails.tokenCount, usageMetadata.candidatesTokenCount)
   }
 
   func testGenerateContent_appCheckNotConfigured_shouldFail() async throws {
@@ -113,6 +127,10 @@ final class IntegrationTests: XCTestCase {
 
     XCTAssertEqual(response.totalTokens, 14)
     XCTAssertEqual(response.totalBillableCharacters, 51)
+    XCTAssertEqual(response.promptTokensDetails.count, 1)
+    let promptTokensDetails = try XCTUnwrap(response.promptTokensDetails.first)
+    XCTAssertEqual(promptTokensDetails.modality, .text)
+    XCTAssertEqual(promptTokensDetails.tokenCount, 14)
   }
 
   #if canImport(UIKit)
@@ -126,6 +144,15 @@ final class IntegrationTests: XCTestCase {
 
       XCTAssertEqual(response.totalTokens, 266)
       XCTAssertEqual(response.totalBillableCharacters, 35)
+      XCTAssertEqual(response.promptTokensDetails.count, 2) // Image prompt + system instruction
+      let textPromptTokensDetails = try XCTUnwrap(response.promptTokensDetails.first {
+        $0.modality == .text
+      }) // System instruction
+      XCTAssertEqual(textPromptTokensDetails.tokenCount, 8)
+      let imagePromptTokenDetails = try XCTUnwrap(response.promptTokensDetails.first {
+        $0.modality == .image
+      })
+      XCTAssertEqual(imagePromptTokenDetails.tokenCount, 258)
     }
   #endif // canImport(UIKit)
 
@@ -137,6 +164,15 @@ final class IntegrationTests: XCTestCase {
 
     XCTAssertEqual(response.totalTokens, 266)
     XCTAssertEqual(response.totalBillableCharacters, 35)
+    XCTAssertEqual(response.promptTokensDetails.count, 2) // Image prompt + system instruction
+    let textPromptTokensDetails = try XCTUnwrap(response.promptTokensDetails.first {
+      $0.modality == .text
+    }) // System instruction
+    XCTAssertEqual(textPromptTokensDetails.tokenCount, 8)
+    let imagePromptTokenDetails = try XCTUnwrap(response.promptTokensDetails.first {
+      $0.modality == .image
+    })
+    XCTAssertEqual(imagePromptTokenDetails.tokenCount, 258)
   }
 
   func testCountTokens_image_fileData_requiresAuth_signedIn() async throws {
@@ -199,6 +235,10 @@ final class IntegrationTests: XCTestCase {
 
     XCTAssertEqual(response.totalTokens, 24)
     XCTAssertEqual(response.totalBillableCharacters, 71)
+    XCTAssertEqual(response.promptTokensDetails.count, 1)
+    let promptTokensDetails = try XCTUnwrap(response.promptTokensDetails.first)
+    XCTAssertEqual(promptTokensDetails.modality, .text)
+    XCTAssertEqual(promptTokensDetails.tokenCount, 24)
   }
 
   func testCountTokens_jsonSchema() async throws {
@@ -220,6 +260,10 @@ final class IntegrationTests: XCTestCase {
 
     XCTAssertEqual(response.totalTokens, 58)
     XCTAssertEqual(response.totalBillableCharacters, 160)
+    XCTAssertEqual(response.promptTokensDetails.count, 1)
+    let promptTokensDetails = try XCTUnwrap(response.promptTokensDetails.first)
+    XCTAssertEqual(promptTokensDetails.modality, .text)
+    XCTAssertEqual(promptTokensDetails.tokenCount, 58)
   }
 
   func testCountTokens_appCheckNotConfigured_shouldFail() async throws {