Fix cache token accounting and tiered pricing

This commit is contained in:
elky
2026-07-10 15:13:12 +08:00
parent 736fc76345
commit 4bf5d4c044
19 changed files with 506 additions and 264 deletions
@@ -949,19 +949,22 @@ fn usage_api_family(api_format: Option<&str>) -> UsageApiFamily {
fn normalize_usage_input_tokens(
api_format: Option<&str>,
input_tokens: i64,
cache_creation_tokens: i64,
cache_read_tokens: i64,
) -> i64 {
if input_tokens <= 0 {
return input_tokens.max(0);
}
if cache_read_tokens <= 0 {
if cache_creation_tokens <= 0 && cache_read_tokens <= 0 {
return input_tokens;
}
match usage_api_family(api_format) {
UsageApiFamily::OpenAi | UsageApiFamily::Gemini => {
(input_tokens - cache_read_tokens).max(0)
}
UsageApiFamily::OpenAi => input_tokens
.saturating_sub(cache_creation_tokens.max(0))
.saturating_sub(cache_read_tokens.max(0))
.max(0),
UsageApiFamily::Gemini => (input_tokens - cache_read_tokens).max(0),
UsageApiFamily::Claude | UsageApiFamily::Unknown => input_tokens,
}
}
@@ -980,9 +983,17 @@ fn normalize_usage_total_input_context(
UsageApiFamily::Claude => {
normalized_input_tokens.saturating_add(normalized_cache_creation_tokens)
}
UsageApiFamily::OpenAi | UsageApiFamily::Gemini => normalize_usage_input_tokens(
UsageApiFamily::OpenAi => normalize_usage_input_tokens(
api_format,
normalized_input_tokens,
normalized_cache_creation_tokens,
normalized_cache_read_tokens,
)
.saturating_add(normalized_cache_creation_tokens),
UsageApiFamily::Gemini => normalize_usage_input_tokens(
api_format,
normalized_input_tokens,
0,
normalized_cache_read_tokens,
),
UsageApiFamily::Unknown => {
@@ -1020,8 +1031,15 @@ fn usage_effective_input_tokens(item: &StoredRequestUsageAudit) -> u64 {
.as_deref()
.or(item.api_format.as_deref());
let input_tokens = i64::try_from(item.input_tokens).unwrap_or(i64::MAX);
let cache_creation_tokens =
i64::try_from(usage_cache_creation_tokens(item)).unwrap_or(i64::MAX);
let cache_read_tokens = i64::try_from(item.cache_read_input_tokens).unwrap_or(i64::MAX);
normalize_usage_input_tokens(api_format, input_tokens, cache_read_tokens) as u64
normalize_usage_input_tokens(
api_format,
input_tokens,
cache_creation_tokens,
cache_read_tokens,
) as u64
}
fn usage_total_tokens(item: &StoredRequestUsageAudit) -> u64 {
@@ -4988,9 +4988,13 @@ normalized_usage AS (
successful_response_time_samples,
CASE
WHEN input_tokens <= 0 THEN 0
WHEN cache_read_tokens <= 0 THEN input_tokens
WHEN split_part(lower(COALESCE(normalized_api_format, '')), ':', 1)
IN ('openai', 'gemini', 'google')
= 'openai'
AND (cache_creation_tokens > 0 OR cache_read_tokens > 0)
THEN GREATEST(input_tokens - cache_creation_tokens - cache_read_tokens, 0)
WHEN split_part(lower(COALESCE(normalized_api_format, '')), ':', 1)
IN ('gemini', 'google')
AND cache_read_tokens > 0
THEN GREATEST(input_tokens - cache_read_tokens, 0)
ELSE input_tokens
END AS effective_input_tokens,
@@ -4999,13 +5003,21 @@ normalized_usage AS (
IN ('claude', 'anthropic')
THEN input_tokens + cache_creation_tokens + cache_read_tokens
WHEN split_part(lower(COALESCE(normalized_api_format, '')), ':', 1)
IN ('openai', 'gemini', 'google')
= 'openai'
THEN (
CASE
WHEN input_tokens <= 0 THEN 0
WHEN cache_read_tokens <= 0 THEN input_tokens
WHEN split_part(lower(COALESCE(normalized_api_format, '')), ':', 1)
IN ('openai', 'gemini', 'google')
WHEN cache_creation_tokens > 0 OR cache_read_tokens > 0
THEN GREATEST(input_tokens - cache_creation_tokens - cache_read_tokens, 0)
ELSE input_tokens
END
) + cache_creation_tokens + cache_read_tokens
WHEN split_part(lower(COALESCE(normalized_api_format, '')), ':', 1)
IN ('gemini', 'google')
THEN (
CASE
WHEN input_tokens <= 0 THEN 0
WHEN cache_read_tokens > 0
THEN GREATEST(input_tokens - cache_read_tokens, 0)
ELSE input_tokens
END
@@ -7210,9 +7222,13 @@ normalized_usage AS (
success_flag,
CASE
WHEN input_tokens <= 0 THEN 0
WHEN cache_read_tokens <= 0 THEN input_tokens
WHEN split_part(lower(COALESCE(normalized_api_format, '')), ':', 1)
IN ('openai', 'gemini', 'google')
= 'openai'
AND (cache_creation_tokens > 0 OR cache_read_tokens > 0)
THEN GREATEST(input_tokens - cache_creation_tokens - cache_read_tokens, 0)
WHEN split_part(lower(COALESCE(normalized_api_format, '')), ':', 1)
IN ('gemini', 'google')
AND cache_read_tokens > 0
THEN GREATEST(input_tokens - cache_read_tokens, 0)
ELSE input_tokens
END AS effective_input_tokens,
@@ -7221,13 +7237,21 @@ normalized_usage AS (
IN ('claude', 'anthropic')
THEN input_tokens + cache_creation_tokens + cache_read_tokens
WHEN split_part(lower(COALESCE(normalized_api_format, '')), ':', 1)
IN ('openai', 'gemini', 'google')
= 'openai'
THEN (
CASE
WHEN input_tokens <= 0 THEN 0
WHEN cache_read_tokens <= 0 THEN input_tokens
WHEN split_part(lower(COALESCE(normalized_api_format, '')), ':', 1)
IN ('openai', 'gemini', 'google')
WHEN cache_creation_tokens > 0 OR cache_read_tokens > 0
THEN GREATEST(input_tokens - cache_creation_tokens - cache_read_tokens, 0)
ELSE input_tokens
END
) + cache_creation_tokens + cache_read_tokens
WHEN split_part(lower(COALESCE(normalized_api_format, '')), ':', 1)
IN ('gemini', 'google')
THEN (
CASE
WHEN input_tokens <= 0 THEN 0
WHEN cache_read_tokens > 0
THEN GREATEST(input_tokens - cache_read_tokens, 0)
ELSE input_tokens
END
@@ -11037,16 +11061,24 @@ fn usage_normalized_api_family(usage: &UpsertUsageRecord) -> String {
fn usage_effective_input_tokens(
input_tokens: Option<i64>,
cache_creation_tokens: Option<i64>,
cache_read_tokens: Option<i64>,
api_family: &str,
) -> Option<i64> {
let input_tokens = input_tokens?;
let cache_creation_tokens = cache_creation_tokens.unwrap_or_default();
let cache_read_tokens = cache_read_tokens.unwrap_or_default();
if matches!(api_family, "openai" | "gemini" | "google")
&& input_tokens > 0
&& cache_read_tokens > 0
{
return Some(input_tokens.saturating_sub(cache_read_tokens));
if input_tokens > 0 {
if api_family == "openai" && (cache_creation_tokens > 0 || cache_read_tokens > 0) {
return Some(
input_tokens
.saturating_sub(cache_creation_tokens)
.saturating_sub(cache_read_tokens),
);
}
if matches!(api_family, "gemini" | "google") && cache_read_tokens > 0 {
return Some(input_tokens.saturating_sub(cache_read_tokens));
}
}
Some(input_tokens)
}
@@ -11145,6 +11177,7 @@ fn usage_settlement_pricing_snapshot_from_usage(
.or_else(|| {
usage_effective_input_tokens(
billing_input_tokens,
billing_cache_creation_tokens,
billing_cache_read_tokens,
api_family.as_str(),
)
@@ -328,7 +328,20 @@ CASE
WHEN (
LOWER(COALESCE(endpoint_api_format, api_format, '')) = 'openai'
OR LOWER(COALESCE(endpoint_api_format, api_format, '')) LIKE 'openai:%'
OR LOWER(COALESCE(endpoint_api_format, api_format, '')) = 'gemini'
)
AND COALESCE(input_tokens, 0) > 0
AND (
COALESCE(cache_creation_input_tokens, 0) > 0
OR COALESCE(cache_read_input_tokens, 0) > 0
)
THEN MAX(
COALESCE(input_tokens, 0)
- COALESCE(cache_creation_input_tokens, 0)
- COALESCE(cache_read_input_tokens, 0),
0
)
WHEN (
LOWER(COALESCE(endpoint_api_format, api_format, '')) = 'gemini'
OR LOWER(COALESCE(endpoint_api_format, api_format, '')) LIKE 'gemini:%'
OR LOWER(COALESCE(endpoint_api_format, api_format, '')) = 'google'
OR LOWER(COALESCE(endpoint_api_format, api_format, '')) LIKE 'google:%'
@@ -345,7 +358,26 @@ CASE
WHEN (
LOWER(COALESCE(endpoint_api_format, api_format, '')) = 'openai'
OR LOWER(COALESCE(endpoint_api_format, api_format, '')) LIKE 'openai:%'
OR LOWER(COALESCE(endpoint_api_format, api_format, '')) = 'gemini'
)
THEN (
CASE
WHEN COALESCE(input_tokens, 0) > 0
AND (
COALESCE(cache_creation_input_tokens, 0) > 0
OR COALESCE(cache_read_input_tokens, 0) > 0
)
THEN MAX(
COALESCE(input_tokens, 0)
- COALESCE(cache_creation_input_tokens, 0)
- COALESCE(cache_read_input_tokens, 0),
0
)
ELSE MAX(COALESCE(input_tokens, 0), 0)
END
) + MAX(COALESCE(cache_creation_input_tokens, 0), 0)
+ MAX(COALESCE(cache_read_input_tokens, 0), 0)
WHEN (
LOWER(COALESCE(endpoint_api_format, api_format, '')) = 'gemini'
OR LOWER(COALESCE(endpoint_api_format, api_format, '')) LIKE 'gemini:%'
OR LOWER(COALESCE(endpoint_api_format, api_format, '')) = 'google'
OR LOWER(COALESCE(endpoint_api_format, api_format, '')) LIKE 'google:%'