diff --git a/FirebaseAI/Sources/GenerateContentResponse.swift b/FirebaseAI/Sources/GenerateContentResponse.swift index cb212e5a616..3304b052bf9 100644 --- a/FirebaseAI/Sources/GenerateContentResponse.swift +++ b/FirebaseAI/Sources/GenerateContentResponse.swift @@ -23,6 +23,9 @@ public struct GenerateContentResponse: Sendable { /// The number of tokens in the request prompt. public let promptTokenCount: Int + /// Number of tokens in the cached part of the prompt (the cached content) + public let cachedContentTokenCount: Int + /// The total number of tokens across the generated response candidates. public let candidatesTokenCount: Int @@ -42,6 +45,9 @@ public struct GenerateContentResponse: Sendable { /// The breakdown, by modality, of how many tokens are consumed by the prompt public let promptTokensDetails: [ModalityTokenCount] + /// The breakdown, by modality, of how many tokens are consumed by the cachedContent + public let cacheTokensDetails: [ModalityTokenCount] + /// The breakdown, by modality, of how many tokens are consumed by the candidates public let candidatesTokensDetails: [ModalityTokenCount] } @@ -339,22 +345,30 @@ extension GenerateContentResponse: Decodable { extension GenerateContentResponse.UsageMetadata: Decodable { enum CodingKeys: CodingKey { case promptTokenCount + case cacheContentTokenCount case candidatesTokenCount case thoughtsTokenCount case totalTokenCount case promptTokensDetails + case cacheTokensDetails case candidatesTokensDetails } public init(from decoder: any Decoder) throws { let container = try decoder.container(keyedBy: CodingKeys.self) promptTokenCount = try container.decodeIfPresent(Int.self, forKey: .promptTokenCount) ?? 0 + cachedContentTokenCount = try container.decodeIfPresent( + Int.self, + forKey: .cacheContentTokenCount + ) ?? 0 candidatesTokenCount = try container.decodeIfPresent(Int.self, forKey: .candidatesTokenCount) ?? 0 thoughtsTokenCount = try container.decodeIfPresent(Int.self, forKey: .thoughtsTokenCount) ?? 0 totalTokenCount = try container.decodeIfPresent(Int.self, forKey: .totalTokenCount) ?? 0 promptTokensDetails = try container.decodeIfPresent([ModalityTokenCount].self, forKey: .promptTokensDetails) ?? [] + cacheTokensDetails = + try container.decodeIfPresent([ModalityTokenCount].self, forKey: .cacheTokensDetails) ?? [] candidatesTokensDetails = try container.decodeIfPresent( [ModalityTokenCount].self, forKey: .candidatesTokensDetails diff --git a/FirebaseAI/Tests/Unit/APITests.swift b/FirebaseAI/Tests/Unit/APITests.swift index 16c963b1f0c..c8c695f1fe2 100644 --- a/FirebaseAI/Tests/Unit/APITests.swift +++ b/FirebaseAI/Tests/Unit/APITests.swift @@ -176,6 +176,7 @@ final class APITests: XCTestCase { // Usage Metadata guard let usageMetadata = response.usageMetadata else { fatalError() } let _: Int = usageMetadata.promptTokenCount + let _: Int = usageMetadata.cachedContentTokenCount let _: Int = usageMetadata.candidatesTokenCount let _: Int = usageMetadata.totalTokenCount