Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Feb 18, 2025
Merged

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

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
ebd0403
[Vertex AI] Add integration tests for modality usage metrics
andrewheard Feb 11, 2025
b233adc
Add `accuracy` threshold for counts to account for whitespace tokens
andrewheard Feb 11, 2025
6e92977
Merge branch 'main' into ah/vertex-modality-integration
andrewheard Feb 11, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
44 changes: 44 additions & 0 deletions FirebaseVertexAI/Tests/TestApp/Tests/Integration/IntegrationTests.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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!
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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)
Expand All @@ -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)

Expand All @@ -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 {
Expand Down Expand Up @@ -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 {
Expand All @@ -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 {
Expand Down
Loading