Skip to content

[FirebaseAI] Implicit caching support #14944

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.

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

Draft
wants to merge 2 commits into
base: pb-rename-to-FirebaseAITestApp
Choose a base branch
from
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
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
14 changes: 14 additions & 0 deletions FirebaseAI/Sources/GenerateContentResponse.swift
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -32,6 +35,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]
}
Expand Down Expand Up @@ -329,20 +335,28 @@ extension GenerateContentResponse: Decodable {
extension GenerateContentResponse.UsageMetadata: Decodable {
enum CodingKeys: CodingKey {
case promptTokenCount
case cacheContentTokenCount
case candidatesTokenCount
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
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
Expand Down
1 change: 1 addition & 0 deletions FirebaseAI/Tests/Unit/APITests.swift
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
Loading