From 0436bea81495f5a4db5a6830bf96f4433f226387 Mon Sep 17 00:00:00 2001 From: Paul Beusterien Date: Fri, 6 Jun 2025 16:06:35 -0700 Subject: [PATCH 1/2] [FirebaseAI] Implicit caching support --- FirebaseAI/Sources/GenerateContentResponse.swift | 11 +++++++++++ FirebaseAI/Tests/Unit/APITests.swift | 1 + 2 files changed, 12 insertions(+) diff --git a/FirebaseAI/Sources/GenerateContentResponse.swift b/FirebaseAI/Sources/GenerateContentResponse.swift index cb212e5a616..66cbd56ba41 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,27 @@ 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 From abb99fc998942527f6757752f65e63408c6a4484 Mon Sep 17 00:00:00 2001 From: Paul Beusterien Date: Fri, 6 Jun 2025 16:20:29 -0700 Subject: [PATCH 2/2] style --- FirebaseAI/Sources/GenerateContentResponse.swift | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/FirebaseAI/Sources/GenerateContentResponse.swift b/FirebaseAI/Sources/GenerateContentResponse.swift index 66cbd56ba41..3304b052bf9 100644 --- a/FirebaseAI/Sources/GenerateContentResponse.swift +++ b/FirebaseAI/Sources/GenerateContentResponse.swift @@ -357,7 +357,10 @@ extension GenerateContentResponse.UsageMetadata: Decodable { 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 + 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