Skip to content

Commit 2e1a8f7

Browse files
authored
[Vertex AI] Add integration tests for modality usage metrics (#14446)
1 parent 4ed50a9 commit 2e1a8f7

File tree

1 file changed

+44
-0
lines changed

1 file changed

+44
-0
lines changed

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

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,8 @@ final class IntegrationTests: XCTestCase {
3939
SafetySetting(harmCategory: .dangerousContent, threshold: .blockLowAndAbove),
4040
SafetySetting(harmCategory: .civicIntegrity, threshold: .blockLowAndAbove),
4141
]
42+
// Candidates and total token counts may differ slightly between runs due to whitespace tokens.
43+
let tokenCountAccuracy = 1
4244

4345
var vertex: VertexAI!
4446
var model: GenerativeModel!
@@ -74,6 +76,18 @@ final class IntegrationTests: XCTestCase {
7476

7577
let text = try XCTUnwrap(response.text).trimmingCharacters(in: .whitespacesAndNewlines)
7678
XCTAssertEqual(text, "Mountain View")
79+
let usageMetadata = try XCTUnwrap(response.usageMetadata)
80+
XCTAssertEqual(usageMetadata.promptTokenCount, 21)
81+
XCTAssertEqual(usageMetadata.candidatesTokenCount, 3, accuracy: tokenCountAccuracy)
82+
XCTAssertEqual(usageMetadata.totalTokenCount, 24, accuracy: tokenCountAccuracy)
83+
XCTAssertEqual(usageMetadata.promptTokensDetails.count, 1)
84+
let promptTokensDetails = try XCTUnwrap(usageMetadata.promptTokensDetails.first)
85+
XCTAssertEqual(promptTokensDetails.modality, .text)
86+
XCTAssertEqual(promptTokensDetails.tokenCount, usageMetadata.promptTokenCount)
87+
XCTAssertEqual(usageMetadata.candidatesTokensDetails.count, 1)
88+
let candidatesTokensDetails = try XCTUnwrap(usageMetadata.candidatesTokensDetails.first)
89+
XCTAssertEqual(candidatesTokensDetails.modality, .text)
90+
XCTAssertEqual(candidatesTokensDetails.tokenCount, usageMetadata.candidatesTokenCount)
7791
}
7892

7993
func testGenerateContent_appCheckNotConfigured_shouldFail() async throws {
@@ -113,6 +127,10 @@ final class IntegrationTests: XCTestCase {
113127

114128
XCTAssertEqual(response.totalTokens, 14)
115129
XCTAssertEqual(response.totalBillableCharacters, 51)
130+
XCTAssertEqual(response.promptTokensDetails.count, 1)
131+
let promptTokensDetails = try XCTUnwrap(response.promptTokensDetails.first)
132+
XCTAssertEqual(promptTokensDetails.modality, .text)
133+
XCTAssertEqual(promptTokensDetails.tokenCount, 14)
116134
}
117135

118136
#if canImport(UIKit)
@@ -126,6 +144,15 @@ final class IntegrationTests: XCTestCase {
126144

127145
XCTAssertEqual(response.totalTokens, 266)
128146
XCTAssertEqual(response.totalBillableCharacters, 35)
147+
XCTAssertEqual(response.promptTokensDetails.count, 2) // Image prompt + system instruction
148+
let textPromptTokensDetails = try XCTUnwrap(response.promptTokensDetails.first {
149+
$0.modality == .text
150+
}) // System instruction
151+
XCTAssertEqual(textPromptTokensDetails.tokenCount, 8)
152+
let imagePromptTokenDetails = try XCTUnwrap(response.promptTokensDetails.first {
153+
$0.modality == .image
154+
})
155+
XCTAssertEqual(imagePromptTokenDetails.tokenCount, 258)
129156
}
130157
#endif // canImport(UIKit)
131158

@@ -137,6 +164,15 @@ final class IntegrationTests: XCTestCase {
137164

138165
XCTAssertEqual(response.totalTokens, 266)
139166
XCTAssertEqual(response.totalBillableCharacters, 35)
167+
XCTAssertEqual(response.promptTokensDetails.count, 2) // Image prompt + system instruction
168+
let textPromptTokensDetails = try XCTUnwrap(response.promptTokensDetails.first {
169+
$0.modality == .text
170+
}) // System instruction
171+
XCTAssertEqual(textPromptTokensDetails.tokenCount, 8)
172+
let imagePromptTokenDetails = try XCTUnwrap(response.promptTokensDetails.first {
173+
$0.modality == .image
174+
})
175+
XCTAssertEqual(imagePromptTokenDetails.tokenCount, 258)
140176
}
141177

142178
func testCountTokens_image_fileData_requiresAuth_signedIn() async throws {
@@ -199,6 +235,10 @@ final class IntegrationTests: XCTestCase {
199235

200236
XCTAssertEqual(response.totalTokens, 24)
201237
XCTAssertEqual(response.totalBillableCharacters, 71)
238+
XCTAssertEqual(response.promptTokensDetails.count, 1)
239+
let promptTokensDetails = try XCTUnwrap(response.promptTokensDetails.first)
240+
XCTAssertEqual(promptTokensDetails.modality, .text)
241+
XCTAssertEqual(promptTokensDetails.tokenCount, 24)
202242
}
203243

204244
func testCountTokens_jsonSchema() async throws {
@@ -220,6 +260,10 @@ final class IntegrationTests: XCTestCase {
220260

221261
XCTAssertEqual(response.totalTokens, 58)
222262
XCTAssertEqual(response.totalBillableCharacters, 160)
263+
XCTAssertEqual(response.promptTokensDetails.count, 1)
264+
let promptTokensDetails = try XCTUnwrap(response.promptTokensDetails.first)
265+
XCTAssertEqual(promptTokensDetails.modality, .text)
266+
XCTAssertEqual(promptTokensDetails.tokenCount, 58)
223267
}
224268

225269
func testCountTokens_appCheckNotConfigured_shouldFail() async throws {

0 commit comments

Comments
 (0)