Skip to content

Commit 1deb75c

Browse files
authored
[Firebase AI] Update Gemma candidate token count integration tests (#14855)
1 parent dbecea0 commit 1deb75c

File tree

1 file changed

+8
-4
lines changed

1 file changed

+8
-4
lines changed

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

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -75,21 +75,25 @@ struct GenerateContentIntegrationTests {
7575

7676
let usageMetadata = try #require(response.usageMetadata)
7777
#expect(usageMetadata.promptTokenCount.isEqual(to: 13, accuracy: tokenCountAccuracy))
78-
#expect(usageMetadata.candidatesTokenCount.isEqual(to: 3, accuracy: tokenCountAccuracy))
79-
#expect(usageMetadata.totalTokenCount.isEqual(to: 16, accuracy: tokenCountAccuracy))
8078
#expect(usageMetadata.promptTokensDetails.count == 1)
8179
let promptTokensDetails = try #require(usageMetadata.promptTokensDetails.first)
8280
#expect(promptTokensDetails.modality == .text)
8381
#expect(promptTokensDetails.tokenCount == usageMetadata.promptTokenCount)
84-
// The field `candidatesTokensDetails` is not included when using Gemma models.
85-
if modelName == ModelNames.gemma3_4B {
82+
// The fields `candidatesTokenCount` and `candidatesTokensDetails` are not included when using
83+
// Gemma models.
84+
if modelName.hasPrefix("gemma") {
85+
#expect(usageMetadata.candidatesTokenCount == 0)
8686
#expect(usageMetadata.candidatesTokensDetails.isEmpty)
8787
} else {
88+
#expect(usageMetadata.candidatesTokenCount.isEqual(to: 3, accuracy: tokenCountAccuracy))
8889
#expect(usageMetadata.candidatesTokensDetails.count == 1)
8990
let candidatesTokensDetails = try #require(usageMetadata.candidatesTokensDetails.first)
9091
#expect(candidatesTokensDetails.modality == .text)
9192
#expect(candidatesTokensDetails.tokenCount == usageMetadata.candidatesTokenCount)
9293
}
94+
#expect(usageMetadata.totalTokenCount > 0)
95+
#expect(usageMetadata.totalTokenCount ==
96+
(usageMetadata.promptTokenCount + usageMetadata.candidatesTokenCount))
9397
}
9498

9599
@Test(

0 commit comments

Comments
 (0)