mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-07 18:07:47 +08:00
elky
778cfb1a5c
Align MySQL and SQLite schemas, migrations, usage, stats, export, and backfill behavior with the shared data contracts. Extend gateway startup and maintenance support across all SQL drivers.
1874 lines
73 KiB
Rust
1874 lines
73 KiB
Rust
use std::collections::{BTreeMap, HashSet};
|
|
use std::time::{SystemTime, UNIX_EPOCH};
|
|
|
|
use aether_ai_formats::UPSTREAM_IS_STREAM_KEY;
|
|
use async_trait::async_trait;
|
|
use sqlx::{mysql::MySqlRow, MySql, QueryBuilder, Row};
|
|
|
|
use aether_data_contracts::repository::usage::{
|
|
strip_deprecated_usage_display_fields, usage_can_recover_terminal_failure,
|
|
usage_request_metadata_client_family, PendingUsageCleanupSummary, StoredRequestUsageAudit,
|
|
StoredUsageDailySummary, StoredUsageDashboardDailyBreakdownRow, StoredUsageDashboardSummary,
|
|
StoredUsageUserTotals, UpsertUsageRecord, UsageCleanupExecutionMode, UsageCleanupPreviewCounts,
|
|
UsageCleanupSummary, UsageCleanupTargets, UsageCleanupWindow, UsageDailyHeatmapQuery,
|
|
UsageDashboardDailyBreakdownQuery, UsageDashboardSummaryQuery, UsageWriteRepository,
|
|
};
|
|
use aether_data_contracts::DataLayerError;
|
|
|
|
use crate::error::SqlResultExt;
|
|
use crate::MysqlPool;
|
|
|
|
mod cleanup;
|
|
mod counters;
|
|
mod http_capture;
|
|
mod read;
|
|
mod snapshots;
|
|
|
|
pub use read::MysqlUsageReadFilter;
|
|
|
|
const USAGE_COLUMNS: &str = r#"
|
|
SELECT
|
|
id,
|
|
`usage`.request_id,
|
|
user_id,
|
|
api_key_id,
|
|
`usage`.username,
|
|
`usage`.api_key_name,
|
|
provider_name,
|
|
model,
|
|
target_model,
|
|
CASE
|
|
WHEN usage_routing_snapshots.request_id IS NOT NULL
|
|
THEN usage_routing_snapshots.selected_provider_id
|
|
ELSE `usage`.provider_id
|
|
END AS provider_id,
|
|
CASE
|
|
WHEN usage_routing_snapshots.request_id IS NOT NULL
|
|
THEN usage_routing_snapshots.selected_endpoint_id
|
|
ELSE `usage`.provider_endpoint_id
|
|
END AS provider_endpoint_id,
|
|
CASE
|
|
WHEN usage_routing_snapshots.request_id IS NOT NULL
|
|
THEN usage_routing_snapshots.selected_provider_api_key_id
|
|
ELSE `usage`.provider_api_key_id
|
|
END AS provider_api_key_id,
|
|
request_type,
|
|
api_format,
|
|
api_family,
|
|
endpoint_kind,
|
|
endpoint_api_format,
|
|
provider_api_family,
|
|
provider_endpoint_kind,
|
|
CASE
|
|
WHEN usage_routing_snapshots.request_id IS NOT NULL
|
|
THEN COALESCE(usage_routing_snapshots.has_format_conversion, FALSE)
|
|
ELSE COALESCE(`usage`.has_format_conversion, FALSE)
|
|
END AS has_format_conversion,
|
|
is_stream,
|
|
upstream_is_stream,
|
|
input_tokens,
|
|
COALESCE(usage_settlement_snapshots.billing_output_tokens, `usage`.output_tokens, 0)
|
|
AS output_tokens,
|
|
total_tokens,
|
|
COALESCE(
|
|
usage_settlement_snapshots.billing_cache_creation_tokens,
|
|
CASE
|
|
WHEN usage_settlement_snapshots.billing_cache_creation_5m_tokens IS NOT NULL
|
|
OR usage_settlement_snapshots.billing_cache_creation_1h_tokens IS NOT NULL
|
|
THEN COALESCE(usage_settlement_snapshots.billing_cache_creation_5m_tokens, 0)
|
|
+ COALESCE(usage_settlement_snapshots.billing_cache_creation_1h_tokens, 0)
|
|
END,
|
|
`usage`.cache_creation_input_tokens,
|
|
0
|
|
) AS cache_creation_input_tokens,
|
|
COALESCE(
|
|
usage_settlement_snapshots.billing_cache_creation_5m_tokens,
|
|
`usage`.cache_creation_ephemeral_5m_input_tokens,
|
|
0
|
|
) AS cache_creation_ephemeral_5m_input_tokens,
|
|
COALESCE(
|
|
usage_settlement_snapshots.billing_cache_creation_1h_tokens,
|
|
`usage`.cache_creation_ephemeral_1h_input_tokens,
|
|
0
|
|
) AS cache_creation_ephemeral_1h_input_tokens,
|
|
COALESCE(
|
|
usage_settlement_snapshots.billing_cache_read_tokens,
|
|
`usage`.cache_read_input_tokens,
|
|
0
|
|
) AS cache_read_input_tokens,
|
|
CAST(COALESCE(
|
|
usage_settlement_snapshots.billing_cache_creation_cost_usd,
|
|
`usage`.cache_creation_cost_usd,
|
|
0
|
|
) AS DOUBLE) AS cache_creation_cost_usd,
|
|
CAST(COALESCE(
|
|
usage_settlement_snapshots.billing_cache_read_cost_usd,
|
|
`usage`.cache_read_cost_usd,
|
|
0
|
|
) AS DOUBLE) AS cache_read_cost_usd,
|
|
CAST(COALESCE(
|
|
usage_settlement_snapshots.output_price_per_1m,
|
|
`usage`.output_price_per_1m
|
|
) AS DOUBLE) AS output_price_per_1m,
|
|
CAST(COALESCE(
|
|
usage_settlement_snapshots.billing_total_cost_usd,
|
|
`usage`.total_cost_usd,
|
|
0
|
|
) AS DOUBLE) AS total_cost_usd,
|
|
CAST(COALESCE(
|
|
usage_settlement_snapshots.billing_actual_total_cost_usd,
|
|
`usage`.actual_total_cost_usd,
|
|
0
|
|
) AS DOUBLE) AS actual_total_cost_usd,
|
|
status_code,
|
|
error_message,
|
|
error_category,
|
|
response_time_ms,
|
|
first_byte_time_ms,
|
|
status,
|
|
COALESCE(usage_settlement_snapshots.billing_status, `usage`.billing_status)
|
|
AS billing_status,
|
|
CAST(COALESCE(usage_http_audits.request_headers, `usage`.request_headers) AS CHAR) AS request_headers,
|
|
CAST(`usage`.request_body AS CHAR) AS request_body,
|
|
`usage`.request_body_compressed,
|
|
CAST(COALESCE(
|
|
usage_http_audits.provider_request_headers,
|
|
`usage`.provider_request_headers
|
|
) AS CHAR) AS provider_request_headers,
|
|
CAST(`usage`.provider_request_body AS CHAR) AS provider_request_body,
|
|
`usage`.provider_request_body_compressed,
|
|
CAST(COALESCE(usage_http_audits.response_headers, `usage`.response_headers) AS CHAR) AS response_headers,
|
|
CAST(`usage`.response_body AS CHAR) AS response_body,
|
|
`usage`.response_body_compressed,
|
|
CAST(COALESCE(
|
|
usage_http_audits.client_response_headers,
|
|
`usage`.client_response_headers
|
|
) AS CHAR) AS client_response_headers,
|
|
CAST(`usage`.client_response_body AS CHAR) AS client_response_body,
|
|
`usage`.client_response_body_compressed,
|
|
usage_http_audits.request_body_ref AS http_request_body_ref,
|
|
usage_http_audits.provider_request_body_ref AS http_provider_request_body_ref,
|
|
usage_http_audits.response_body_ref AS http_response_body_ref,
|
|
usage_http_audits.client_response_body_ref AS http_client_response_body_ref,
|
|
usage_http_audits.request_body_state AS http_request_body_state,
|
|
usage_http_audits.provider_request_body_state AS http_provider_request_body_state,
|
|
usage_http_audits.response_body_state AS http_response_body_state,
|
|
usage_http_audits.client_response_body_state AS http_client_response_body_state,
|
|
request_metadata,
|
|
CASE
|
|
WHEN usage_routing_snapshots.request_id IS NOT NULL
|
|
THEN usage_routing_snapshots.candidate_id
|
|
ELSE `usage`.candidate_id
|
|
END AS routing_candidate_id,
|
|
CASE
|
|
WHEN usage_routing_snapshots.request_id IS NOT NULL
|
|
THEN usage_routing_snapshots.candidate_index
|
|
ELSE `usage`.candidate_index
|
|
END AS routing_candidate_index,
|
|
CASE
|
|
WHEN usage_routing_snapshots.request_id IS NOT NULL
|
|
THEN usage_routing_snapshots.key_name
|
|
ELSE `usage`.key_name
|
|
END AS routing_key_name,
|
|
CASE
|
|
WHEN usage_routing_snapshots.request_id IS NOT NULL
|
|
THEN usage_routing_snapshots.planner_kind
|
|
ELSE `usage`.planner_kind
|
|
END AS routing_planner_kind,
|
|
CASE
|
|
WHEN usage_routing_snapshots.request_id IS NOT NULL
|
|
THEN usage_routing_snapshots.route_family
|
|
ELSE `usage`.route_family
|
|
END AS routing_route_family,
|
|
CASE
|
|
WHEN usage_routing_snapshots.request_id IS NOT NULL
|
|
THEN usage_routing_snapshots.route_kind
|
|
ELSE `usage`.route_kind
|
|
END AS routing_route_kind,
|
|
CASE
|
|
WHEN usage_routing_snapshots.request_id IS NOT NULL
|
|
THEN usage_routing_snapshots.execution_path
|
|
ELSE `usage`.execution_path
|
|
END AS routing_execution_path,
|
|
CASE
|
|
WHEN usage_routing_snapshots.request_id IS NOT NULL
|
|
THEN usage_routing_snapshots.local_execution_runtime_miss_reason
|
|
ELSE `usage`.local_execution_runtime_miss_reason
|
|
END AS routing_local_execution_runtime_miss_reason,
|
|
usage_settlement_snapshots.billing_snapshot_schema_version
|
|
AS settlement_billing_snapshot_schema_version,
|
|
usage_settlement_snapshots.billing_snapshot_status AS settlement_billing_snapshot_status,
|
|
CAST(usage_settlement_snapshots.rate_multiplier AS DOUBLE) AS settlement_rate_multiplier,
|
|
usage_settlement_snapshots.is_free_tier AS settlement_is_free_tier,
|
|
CAST(usage_settlement_snapshots.input_price_per_1m AS DOUBLE)
|
|
AS settlement_input_price_per_1m,
|
|
CAST(usage_settlement_snapshots.output_price_per_1m AS DOUBLE)
|
|
AS settlement_output_price_per_1m,
|
|
CAST(usage_settlement_snapshots.cache_creation_price_per_1m AS DOUBLE)
|
|
AS settlement_cache_creation_price_per_1m,
|
|
CAST(usage_settlement_snapshots.cache_read_price_per_1m AS DOUBLE)
|
|
AS settlement_cache_read_price_per_1m,
|
|
CAST(usage_settlement_snapshots.price_per_request AS DOUBLE)
|
|
AS settlement_price_per_request,
|
|
usage_settlement_snapshots.settlement_snapshot_schema_version
|
|
AS settlement_snapshot_schema_version,
|
|
CAST(usage_settlement_snapshots.settlement_snapshot AS CHAR) AS settlement_snapshot,
|
|
CAST(usage_settlement_snapshots.billing_dimensions AS CHAR)
|
|
AS settlement_billing_dimensions,
|
|
usage_settlement_snapshots.billing_input_tokens AS settlement_billing_input_tokens,
|
|
usage_settlement_snapshots.billing_effective_input_tokens
|
|
AS settlement_billing_effective_input_tokens,
|
|
usage_settlement_snapshots.billing_output_tokens AS settlement_billing_output_tokens,
|
|
usage_settlement_snapshots.billing_cache_creation_tokens
|
|
AS settlement_billing_cache_creation_tokens,
|
|
usage_settlement_snapshots.billing_cache_creation_5m_tokens
|
|
AS settlement_billing_cache_creation_5m_tokens,
|
|
usage_settlement_snapshots.billing_cache_creation_1h_tokens
|
|
AS settlement_billing_cache_creation_1h_tokens,
|
|
usage_settlement_snapshots.billing_cache_read_tokens
|
|
AS settlement_billing_cache_read_tokens,
|
|
usage_settlement_snapshots.billing_total_input_context
|
|
AS settlement_billing_total_input_context,
|
|
CAST(usage_settlement_snapshots.billing_cache_creation_cost_usd AS DOUBLE)
|
|
AS settlement_billing_cache_creation_cost_usd,
|
|
CAST(usage_settlement_snapshots.billing_cache_read_cost_usd AS DOUBLE)
|
|
AS settlement_billing_cache_read_cost_usd,
|
|
CAST(usage_settlement_snapshots.billing_total_cost_usd AS DOUBLE)
|
|
AS settlement_billing_total_cost_usd,
|
|
CAST(usage_settlement_snapshots.billing_actual_total_cost_usd AS DOUBLE)
|
|
AS settlement_billing_actual_total_cost_usd,
|
|
usage_settlement_snapshots.billing_pricing_source AS settlement_billing_pricing_source,
|
|
usage_settlement_snapshots.billing_rule_id AS settlement_billing_rule_id,
|
|
usage_settlement_snapshots.billing_rule_version AS settlement_billing_rule_version,
|
|
COALESCE(usage_settlement_snapshots.finalized_at, `usage`.finalized_at)
|
|
AS finalized_at_unix_secs,
|
|
created_at_unix_ms,
|
|
updated_at_unix_secs
|
|
FROM `usage`
|
|
LEFT JOIN usage_http_audits
|
|
ON usage_http_audits.request_id = `usage`.request_id
|
|
LEFT JOIN usage_routing_snapshots
|
|
ON usage_routing_snapshots.request_id = `usage`.request_id
|
|
LEFT JOIN usage_settlement_snapshots
|
|
ON usage_settlement_snapshots.request_id = `usage`.request_id
|
|
"#;
|
|
|
|
const UPSERT_USAGE_SQL: &str = r#"
|
|
INSERT INTO `usage` (
|
|
request_id,
|
|
id,
|
|
user_id,
|
|
api_key_id,
|
|
provider_name,
|
|
model,
|
|
target_model,
|
|
provider_id,
|
|
provider_endpoint_id,
|
|
provider_api_key_id,
|
|
request_type,
|
|
api_format,
|
|
api_family,
|
|
endpoint_kind,
|
|
endpoint_api_format,
|
|
provider_api_family,
|
|
provider_endpoint_kind,
|
|
has_format_conversion,
|
|
is_stream,
|
|
upstream_is_stream,
|
|
input_tokens,
|
|
output_tokens,
|
|
total_tokens,
|
|
cache_creation_input_tokens,
|
|
cache_creation_ephemeral_5m_input_tokens,
|
|
cache_creation_ephemeral_1h_input_tokens,
|
|
cache_read_input_tokens,
|
|
cache_creation_cost_usd,
|
|
cache_read_cost_usd,
|
|
output_price_per_1m,
|
|
total_cost_usd,
|
|
actual_total_cost_usd,
|
|
status_code,
|
|
error_message,
|
|
error_category,
|
|
response_time_ms,
|
|
first_byte_time_ms,
|
|
status,
|
|
billing_status,
|
|
request_metadata,
|
|
candidate_id,
|
|
candidate_index,
|
|
key_name,
|
|
planner_kind,
|
|
route_family,
|
|
route_kind,
|
|
execution_path,
|
|
local_execution_runtime_miss_reason,
|
|
finalized_at,
|
|
created_at_unix_ms,
|
|
updated_at_unix_secs
|
|
) VALUES (
|
|
?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?,
|
|
?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?,
|
|
?
|
|
)
|
|
ON DUPLICATE KEY UPDATE
|
|
user_id = VALUES(user_id),
|
|
api_key_id = VALUES(api_key_id),
|
|
provider_name = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN provider_name ELSE VALUES(provider_name) END,
|
|
model = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN model ELSE VALUES(model) END,
|
|
target_model = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN target_model ELSE VALUES(target_model) END,
|
|
provider_id = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN provider_id ELSE VALUES(provider_id) END,
|
|
provider_endpoint_id = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN provider_endpoint_id ELSE VALUES(provider_endpoint_id) END,
|
|
provider_api_key_id = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN provider_api_key_id ELSE VALUES(provider_api_key_id) END,
|
|
request_type = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN request_type ELSE VALUES(request_type) END,
|
|
api_format = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN api_format ELSE VALUES(api_format) END,
|
|
api_family = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN api_family ELSE VALUES(api_family) END,
|
|
endpoint_kind = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN endpoint_kind ELSE VALUES(endpoint_kind) END,
|
|
endpoint_api_format = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN endpoint_api_format ELSE VALUES(endpoint_api_format) END,
|
|
provider_api_family = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN provider_api_family ELSE VALUES(provider_api_family) END,
|
|
provider_endpoint_kind = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN provider_endpoint_kind ELSE VALUES(provider_endpoint_kind) END,
|
|
has_format_conversion = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN has_format_conversion ELSE VALUES(has_format_conversion) END,
|
|
is_stream = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN is_stream ELSE VALUES(is_stream) END,
|
|
upstream_is_stream = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN upstream_is_stream ELSE VALUES(upstream_is_stream) END,
|
|
input_tokens = CASE
|
|
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN input_tokens
|
|
ELSE VALUES(input_tokens)
|
|
END,
|
|
output_tokens = CASE
|
|
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN output_tokens
|
|
ELSE VALUES(output_tokens)
|
|
END,
|
|
total_tokens = CASE
|
|
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN total_tokens
|
|
ELSE VALUES(total_tokens)
|
|
END,
|
|
cache_creation_input_tokens = CASE
|
|
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN cache_creation_input_tokens
|
|
ELSE VALUES(cache_creation_input_tokens)
|
|
END,
|
|
cache_creation_ephemeral_5m_input_tokens = CASE
|
|
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN cache_creation_ephemeral_5m_input_tokens
|
|
ELSE VALUES(cache_creation_ephemeral_5m_input_tokens)
|
|
END,
|
|
cache_creation_ephemeral_1h_input_tokens = CASE
|
|
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN cache_creation_ephemeral_1h_input_tokens
|
|
ELSE VALUES(cache_creation_ephemeral_1h_input_tokens)
|
|
END,
|
|
cache_read_input_tokens = CASE
|
|
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN cache_read_input_tokens
|
|
ELSE VALUES(cache_read_input_tokens)
|
|
END,
|
|
cache_creation_cost_usd = CASE
|
|
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN cache_creation_cost_usd
|
|
ELSE VALUES(cache_creation_cost_usd)
|
|
END,
|
|
cache_read_cost_usd = CASE
|
|
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN cache_read_cost_usd
|
|
ELSE VALUES(cache_read_cost_usd)
|
|
END,
|
|
output_price_per_1m = CASE
|
|
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN output_price_per_1m
|
|
ELSE VALUES(output_price_per_1m)
|
|
END,
|
|
total_cost_usd = CASE
|
|
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN total_cost_usd
|
|
ELSE VALUES(total_cost_usd)
|
|
END,
|
|
actual_total_cost_usd = CASE
|
|
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN actual_total_cost_usd
|
|
ELSE VALUES(actual_total_cost_usd)
|
|
END,
|
|
status_code = CASE
|
|
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN status_code
|
|
WHEN status = 'streaming' AND VALUES(status) = 'pending' THEN status_code
|
|
WHEN status = 'streaming' AND VALUES(status) = 'streaming' AND VALUES(status_code) IS NULL THEN status_code
|
|
ELSE VALUES(status_code)
|
|
END,
|
|
error_message = CASE
|
|
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN error_message
|
|
WHEN status = 'streaming' AND VALUES(status) = 'pending' THEN error_message
|
|
ELSE VALUES(error_message)
|
|
END,
|
|
error_category = CASE
|
|
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN error_category
|
|
WHEN status = 'streaming' AND VALUES(status) = 'pending' THEN error_category
|
|
ELSE VALUES(error_category)
|
|
END,
|
|
response_time_ms = CASE
|
|
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN response_time_ms
|
|
WHEN VALUES(response_time_ms) IS NULL OR VALUES(response_time_ms) = 0
|
|
THEN COALESCE(response_time_ms, VALUES(response_time_ms))
|
|
ELSE VALUES(response_time_ms)
|
|
END,
|
|
first_byte_time_ms = CASE
|
|
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN first_byte_time_ms
|
|
WHEN VALUES(first_byte_time_ms) IS NULL OR VALUES(first_byte_time_ms) = 0
|
|
THEN COALESCE(first_byte_time_ms, VALUES(first_byte_time_ms))
|
|
ELSE VALUES(first_byte_time_ms)
|
|
END,
|
|
billing_status = CASE
|
|
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN billing_status
|
|
ELSE VALUES(billing_status)
|
|
END,
|
|
request_metadata = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN request_metadata ELSE VALUES(request_metadata) END,
|
|
candidate_id = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN candidate_id ELSE VALUES(candidate_id) END,
|
|
candidate_index = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN candidate_index ELSE VALUES(candidate_index) END,
|
|
key_name = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN key_name ELSE VALUES(key_name) END,
|
|
planner_kind = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN planner_kind ELSE VALUES(planner_kind) END,
|
|
route_family = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN route_family ELSE VALUES(route_family) END,
|
|
route_kind = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN route_kind ELSE VALUES(route_kind) END,
|
|
execution_path = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN execution_path ELSE VALUES(execution_path) END,
|
|
local_execution_runtime_miss_reason = CASE WHEN (status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')) OR (status = 'streaming' AND VALUES(status) = 'pending') THEN local_execution_runtime_miss_reason ELSE VALUES(local_execution_runtime_miss_reason) END,
|
|
finalized_at = CASE
|
|
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN finalized_at
|
|
ELSE VALUES(finalized_at)
|
|
END,
|
|
updated_at_unix_secs = CASE
|
|
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN updated_at_unix_secs
|
|
ELSE VALUES(updated_at_unix_secs)
|
|
END,
|
|
status = CASE
|
|
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN status
|
|
WHEN status = 'streaming' AND VALUES(status) = 'pending' THEN status
|
|
ELSE VALUES(status)
|
|
END
|
|
"#;
|
|
|
|
const SELECT_STALE_PENDING_USAGE_BATCH_SQL: &str = r#"
|
|
SELECT
|
|
`usage`.request_id,
|
|
`usage`.status,
|
|
COALESCE(usage_settlement_snapshots.billing_status, `usage`.billing_status) AS billing_status
|
|
FROM `usage`
|
|
LEFT JOIN usage_settlement_snapshots
|
|
ON usage_settlement_snapshots.request_id = `usage`.request_id
|
|
WHERE `usage`.status IN ('pending', 'streaming')
|
|
AND `usage`.created_at_unix_ms < ?
|
|
ORDER BY `usage`.created_at_unix_ms ASC, `usage`.request_id ASC
|
|
LIMIT ?
|
|
"#;
|
|
|
|
const SELECT_COMPLETED_REQUEST_CANDIDATES_SQL: &str = r#"
|
|
SELECT status, extra_data
|
|
FROM request_candidates
|
|
WHERE request_id = ?
|
|
AND status IN ('streaming', 'success')
|
|
"#;
|
|
|
|
const MYSQL_USAGE_CANONICAL_TOTAL_TOKENS_EXPR: &str = r#"
|
|
GREATEST(
|
|
COALESCE(
|
|
CASE
|
|
WHEN settlement.billing_effective_input_tokens IS NOT NULL THEN
|
|
GREATEST(settlement.billing_effective_input_tokens, 0)
|
|
+ GREATEST(COALESCE(settlement.billing_output_tokens, `usage`.output_tokens, 0), 0)
|
|
+ GREATEST(
|
|
COALESCE(
|
|
settlement.billing_cache_creation_tokens,
|
|
CASE
|
|
WHEN settlement.billing_cache_creation_5m_tokens IS NOT NULL
|
|
OR settlement.billing_cache_creation_1h_tokens IS NOT NULL
|
|
THEN COALESCE(settlement.billing_cache_creation_5m_tokens, 0)
|
|
+ COALESCE(settlement.billing_cache_creation_1h_tokens, 0)
|
|
END,
|
|
CASE
|
|
WHEN COALESCE(`usage`.cache_creation_input_tokens, 0) = 0
|
|
AND (
|
|
COALESCE(`usage`.cache_creation_ephemeral_5m_input_tokens, 0)
|
|
+ COALESCE(`usage`.cache_creation_ephemeral_1h_input_tokens, 0)
|
|
) > 0
|
|
THEN COALESCE(`usage`.cache_creation_ephemeral_5m_input_tokens, 0)
|
|
+ COALESCE(`usage`.cache_creation_ephemeral_1h_input_tokens, 0)
|
|
ELSE COALESCE(`usage`.cache_creation_input_tokens, 0)
|
|
END,
|
|
0
|
|
),
|
|
0
|
|
)
|
|
+ GREATEST(
|
|
COALESCE(
|
|
settlement.billing_cache_read_tokens,
|
|
`usage`.cache_read_input_tokens,
|
|
0
|
|
),
|
|
0
|
|
)
|
|
WHEN settlement.billing_total_input_context IS NOT NULL THEN
|
|
GREATEST(settlement.billing_total_input_context, 0)
|
|
+ GREATEST(COALESCE(settlement.billing_output_tokens, `usage`.output_tokens, 0), 0)
|
|
END,
|
|
NULLIF(GREATEST(COALESCE(`usage`.total_tokens, 0), 0), 0),
|
|
(
|
|
CASE
|
|
WHEN (
|
|
LOWER(COALESCE(`usage`.endpoint_api_format, `usage`.api_format, '')) = 'openai'
|
|
OR LOWER(COALESCE(`usage`.endpoint_api_format, `usage`.api_format, '')) LIKE 'openai:%'
|
|
OR LOWER(COALESCE(`usage`.endpoint_api_format, `usage`.api_format, '')) = 'gemini'
|
|
OR LOWER(COALESCE(`usage`.endpoint_api_format, `usage`.api_format, '')) LIKE 'gemini:%'
|
|
OR LOWER(COALESCE(`usage`.endpoint_api_format, `usage`.api_format, '')) = 'google'
|
|
OR LOWER(COALESCE(`usage`.endpoint_api_format, `usage`.api_format, '')) LIKE 'google:%'
|
|
)
|
|
AND COALESCE(`usage`.input_tokens, 0) > 0
|
|
AND COALESCE(`usage`.cache_read_input_tokens, 0) > 0
|
|
THEN GREATEST(
|
|
COALESCE(`usage`.input_tokens, 0) - COALESCE(`usage`.cache_read_input_tokens, 0),
|
|
0
|
|
)
|
|
ELSE GREATEST(COALESCE(`usage`.input_tokens, 0), 0)
|
|
END
|
|
+ GREATEST(COALESCE(`usage`.output_tokens, 0), 0)
|
|
+ (
|
|
CASE
|
|
WHEN COALESCE(`usage`.cache_creation_input_tokens, 0) = 0
|
|
AND (
|
|
COALESCE(`usage`.cache_creation_ephemeral_5m_input_tokens, 0)
|
|
+ COALESCE(`usage`.cache_creation_ephemeral_1h_input_tokens, 0)
|
|
) > 0
|
|
THEN COALESCE(`usage`.cache_creation_ephemeral_5m_input_tokens, 0)
|
|
+ COALESCE(`usage`.cache_creation_ephemeral_1h_input_tokens, 0)
|
|
ELSE GREATEST(COALESCE(`usage`.cache_creation_input_tokens, 0), 0)
|
|
END
|
|
)
|
|
+ GREATEST(COALESCE(`usage`.cache_read_input_tokens, 0), 0)
|
|
),
|
|
0
|
|
),
|
|
0
|
|
)
|
|
"#;
|
|
|
|
const MYSQL_PROVIDER_KEY_SUCCESS_FLAG_EXPR: &str = r#"
|
|
CASE
|
|
WHEN status IN ('completed', 'success', 'ok', 'billed', 'settled')
|
|
AND (status_code IS NULL OR status_code < 400)
|
|
AND (error_message IS NULL OR TRIM(error_message) = '')
|
|
THEN 1
|
|
ELSE 0
|
|
END
|
|
"#;
|
|
|
|
const MYSQL_PROVIDER_KEY_ERROR_FLAG_EXPR: &str = r#"
|
|
CASE
|
|
WHEN status NOT IN ('pending', 'streaming')
|
|
AND NOT (
|
|
status IN ('completed', 'success', 'ok', 'billed', 'settled')
|
|
AND (status_code IS NULL OR status_code < 400)
|
|
AND (error_message IS NULL OR TRIM(error_message) = '')
|
|
)
|
|
THEN 1
|
|
ELSE 0
|
|
END
|
|
"#;
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub struct MysqlUsageWriteRepository {
|
|
pool: MysqlPool,
|
|
}
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub struct MysqlUsageStorage {
|
|
pool: MysqlPool,
|
|
}
|
|
|
|
impl MysqlUsageStorage {
|
|
pub fn new(pool: MysqlPool) -> Self {
|
|
Self { pool }
|
|
}
|
|
|
|
pub async fn read_usage_counter_health(
|
|
&self,
|
|
) -> Result<aether_data_contracts::repository::usage::UsageCounterHealthSnapshot, DataLayerError>
|
|
{
|
|
counters::read_health(&self.pool).await
|
|
}
|
|
|
|
pub async fn read_usage_counter_pending_health(
|
|
&self,
|
|
) -> Result<
|
|
aether_data_contracts::repository::usage::UsageCounterPendingHealthSnapshot,
|
|
DataLayerError,
|
|
> {
|
|
counters::read_pending_health(&self.pool).await
|
|
}
|
|
|
|
async fn summarize_usage_daily_heatmap_raw_from_range(
|
|
&self,
|
|
created_from_unix_secs: u64,
|
|
created_until_unix_secs: u64,
|
|
user_id: Option<&str>,
|
|
) -> Result<Vec<StoredUsageDailySummary>, DataLayerError> {
|
|
let mut sql = format!(
|
|
r#"
|
|
SELECT
|
|
DATE_FORMAT(FROM_UNIXTIME(created_at_unix_ms), '%Y-%m-%d') AS date,
|
|
CAST(COUNT(*) AS SIGNED) AS requests,
|
|
CAST(COALESCE(SUM({canonical_total_tokens_expr}), 0) AS SIGNED) AS total_tokens,
|
|
CAST(COALESCE(SUM(COALESCE(total_cost_usd, 0)), 0) AS DOUBLE) AS total_cost_usd,
|
|
CAST(COALESCE(SUM(COALESCE(actual_total_cost_usd, 0)), 0) AS DOUBLE) AS actual_total_cost_usd
|
|
FROM `usage`
|
|
LEFT JOIN usage_settlement_snapshots AS settlement
|
|
ON settlement.request_id = `usage`.request_id
|
|
WHERE created_at_unix_ms >= ?
|
|
AND created_at_unix_ms < ?
|
|
AND status NOT IN ('pending', 'streaming')
|
|
AND provider_name NOT IN ('unknown', 'pending')
|
|
"#,
|
|
canonical_total_tokens_expr = MYSQL_USAGE_CANONICAL_TOTAL_TOKENS_EXPR,
|
|
);
|
|
if user_id.is_some() {
|
|
sql.push_str(" AND user_id = ?\n");
|
|
}
|
|
sql.push_str("GROUP BY date ORDER BY date ASC");
|
|
|
|
let mut query = sqlx::query(&sql)
|
|
.bind(to_i64(created_from_unix_secs, "usage.created_at_unix_ms")?)
|
|
.bind(to_i64(created_until_unix_secs, "usage.created_at_unix_ms")?);
|
|
if let Some(user_id) = user_id {
|
|
query = query.bind(user_id.to_string());
|
|
}
|
|
let rows = query.fetch_all(&self.pool).await.map_sql_err()?;
|
|
rows.iter().map(map_mysql_usage_daily_summary).collect()
|
|
}
|
|
|
|
async fn summarize_usage_daily_heatmap_from_daily_aggregates(
|
|
&self,
|
|
created_from_unix_secs: u64,
|
|
created_until_unix_secs: u64,
|
|
user_id: Option<&str>,
|
|
) -> Result<Vec<StoredUsageDailySummary>, DataLayerError> {
|
|
let rows = if let Some(user_id) = user_id {
|
|
sqlx::query(
|
|
r#"
|
|
SELECT
|
|
DATE_FORMAT(FROM_UNIXTIME(`date`), '%Y-%m-%d') AS date,
|
|
total_requests AS requests,
|
|
input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens AS total_tokens,
|
|
total_cost AS total_cost_usd,
|
|
total_cost AS actual_total_cost_usd
|
|
FROM stats_user_daily
|
|
WHERE user_id = ?
|
|
AND `date` >= ?
|
|
AND `date` < ?
|
|
AND total_requests > 0
|
|
ORDER BY `date` ASC
|
|
"#,
|
|
)
|
|
.bind(user_id)
|
|
.bind(to_i64(created_from_unix_secs, "stats_user_daily.date")?)
|
|
.bind(to_i64(created_until_unix_secs, "stats_user_daily.date")?)
|
|
.fetch_all(&self.pool)
|
|
.await
|
|
.map_sql_err()?
|
|
} else {
|
|
sqlx::query(
|
|
r#"
|
|
SELECT
|
|
DATE_FORMAT(FROM_UNIXTIME(`date`), '%Y-%m-%d') AS date,
|
|
total_requests AS requests,
|
|
input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens AS total_tokens,
|
|
total_cost AS total_cost_usd,
|
|
actual_total_cost AS actual_total_cost_usd
|
|
FROM stats_daily
|
|
WHERE `date` >= ?
|
|
AND `date` < ?
|
|
AND total_requests > 0
|
|
ORDER BY `date` ASC
|
|
"#,
|
|
)
|
|
.bind(to_i64(created_from_unix_secs, "stats_daily.date")?)
|
|
.bind(to_i64(created_until_unix_secs, "stats_daily.date")?)
|
|
.fetch_all(&self.pool)
|
|
.await
|
|
.map_sql_err()?
|
|
};
|
|
|
|
rows.iter().map(map_mysql_usage_daily_summary).collect()
|
|
}
|
|
|
|
pub async fn summarize_usage_daily_heatmap(
|
|
&self,
|
|
query: &UsageDailyHeatmapQuery,
|
|
) -> Result<Vec<StoredUsageDailySummary>, DataLayerError> {
|
|
let created_until_unix_secs = usage_current_unix_secs().saturating_add(1);
|
|
let user_id = query.user_id.as_deref();
|
|
let mut summaries = BTreeMap::<String, StoredUsageDailySummary>::new();
|
|
|
|
for item in self
|
|
.summarize_usage_daily_heatmap_from_daily_aggregates(
|
|
query.created_from_unix_secs,
|
|
created_until_unix_secs,
|
|
user_id,
|
|
)
|
|
.await?
|
|
{
|
|
summaries.insert(item.date.clone(), item);
|
|
}
|
|
for item in self
|
|
.summarize_usage_daily_heatmap_raw_from_range(
|
|
query.created_from_unix_secs,
|
|
created_until_unix_secs,
|
|
user_id,
|
|
)
|
|
.await?
|
|
{
|
|
summaries.entry(item.date.clone()).or_insert(item);
|
|
}
|
|
|
|
Ok(summaries.into_values().collect())
|
|
}
|
|
|
|
pub async fn summarize_dashboard_usage_from_daily_aggregates(
|
|
&self,
|
|
query: &UsageDashboardSummaryQuery,
|
|
) -> Result<Option<StoredUsageDashboardSummary>, DataLayerError> {
|
|
let row = if let Some(user_id) = query.user_id.as_deref() {
|
|
sqlx::query(
|
|
r#"
|
|
SELECT
|
|
CAST(COALESCE(SUM(total_requests), 0) AS SIGNED) AS total_requests,
|
|
CAST(COALESCE(SUM(input_tokens), 0) AS SIGNED) AS input_tokens,
|
|
CAST(COALESCE(SUM(input_tokens), 0) AS SIGNED) AS effective_input_tokens,
|
|
CAST(COALESCE(SUM(output_tokens), 0) AS SIGNED) AS output_tokens,
|
|
CAST(COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) AS SIGNED) AS total_tokens,
|
|
CAST(COALESCE(SUM(cache_creation_tokens), 0) AS SIGNED) AS cache_creation_tokens,
|
|
CAST(COALESCE(SUM(cache_read_tokens), 0) AS SIGNED) AS cache_read_tokens,
|
|
CAST(COALESCE(SUM(input_tokens + cache_creation_tokens + cache_read_tokens), 0) AS SIGNED) AS total_input_context,
|
|
CAST(0.0 AS DOUBLE) AS cache_creation_cost_usd,
|
|
CAST(0.0 AS DOUBLE) AS cache_read_cost_usd,
|
|
CAST(COALESCE(SUM(COALESCE(total_cost, 0)), 0) AS DOUBLE) AS total_cost_usd,
|
|
CAST(COALESCE(SUM(COALESCE(total_cost, 0)), 0) AS DOUBLE) AS actual_total_cost_usd,
|
|
CAST(COALESCE(SUM(error_requests), 0) AS SIGNED) AS error_requests,
|
|
CAST(0.0 AS DOUBLE) AS response_time_sum_ms,
|
|
CAST(0 AS SIGNED) AS response_time_samples
|
|
FROM stats_user_daily
|
|
WHERE user_id = ?
|
|
AND `date` >= ?
|
|
AND `date` < ?
|
|
"#,
|
|
)
|
|
.bind(user_id)
|
|
.bind(to_i64(
|
|
query.created_from_unix_secs,
|
|
"stats_user_daily.date",
|
|
)?)
|
|
.bind(to_i64(
|
|
query.created_until_unix_secs,
|
|
"stats_user_daily.date",
|
|
)?)
|
|
.fetch_one(&self.pool)
|
|
.await
|
|
.map_sql_err()?
|
|
} else {
|
|
sqlx::query(
|
|
r#"
|
|
SELECT
|
|
CAST(COALESCE(SUM(total_requests), 0) AS SIGNED) AS total_requests,
|
|
CAST(COALESCE(SUM(input_tokens), 0) AS SIGNED) AS input_tokens,
|
|
CAST(COALESCE(SUM(input_tokens), 0) AS SIGNED) AS effective_input_tokens,
|
|
CAST(COALESCE(SUM(output_tokens), 0) AS SIGNED) AS output_tokens,
|
|
CAST(COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) AS SIGNED) AS total_tokens,
|
|
CAST(COALESCE(SUM(cache_creation_tokens), 0) AS SIGNED) AS cache_creation_tokens,
|
|
CAST(COALESCE(SUM(cache_read_tokens), 0) AS SIGNED) AS cache_read_tokens,
|
|
CAST(COALESCE(SUM(input_tokens + cache_creation_tokens + cache_read_tokens), 0) AS SIGNED) AS total_input_context,
|
|
CAST(COALESCE(SUM(COALESCE(cache_creation_cost, 0)), 0) AS DOUBLE) AS cache_creation_cost_usd,
|
|
CAST(COALESCE(SUM(COALESCE(cache_read_cost, 0)), 0) AS DOUBLE) AS cache_read_cost_usd,
|
|
CAST(COALESCE(SUM(COALESCE(total_cost, 0)), 0) AS DOUBLE) AS total_cost_usd,
|
|
CAST(COALESCE(SUM(COALESCE(actual_total_cost, 0)), 0) AS DOUBLE) AS actual_total_cost_usd,
|
|
CAST(COALESCE(SUM(error_requests), 0) AS SIGNED) AS error_requests,
|
|
CAST(0.0 AS DOUBLE) AS response_time_sum_ms,
|
|
CAST(0 AS SIGNED) AS response_time_samples
|
|
FROM stats_daily
|
|
WHERE `date` >= ?
|
|
AND `date` < ?
|
|
"#,
|
|
)
|
|
.bind(to_i64(query.created_from_unix_secs, "stats_daily.date")?)
|
|
.bind(to_i64(query.created_until_unix_secs, "stats_daily.date")?)
|
|
.fetch_one(&self.pool)
|
|
.await
|
|
.map_sql_err()?
|
|
};
|
|
|
|
let total_requests = row_u64(&row, "total_requests")?;
|
|
if total_requests == 0 {
|
|
return Ok(None);
|
|
}
|
|
|
|
Ok(Some(StoredUsageDashboardSummary {
|
|
total_requests,
|
|
input_tokens: row_u64(&row, "input_tokens")?,
|
|
effective_input_tokens: row_u64(&row, "effective_input_tokens")?,
|
|
output_tokens: row_u64(&row, "output_tokens")?,
|
|
total_tokens: row_u64(&row, "total_tokens")?,
|
|
cache_creation_tokens: row_u64(&row, "cache_creation_tokens")?,
|
|
cache_read_tokens: row_u64(&row, "cache_read_tokens")?,
|
|
total_input_context: row_u64(&row, "total_input_context")?,
|
|
cache_creation_cost_usd: row.try_get("cache_creation_cost_usd").map_sql_err()?,
|
|
cache_read_cost_usd: row.try_get("cache_read_cost_usd").map_sql_err()?,
|
|
total_cost_usd: row.try_get("total_cost_usd").map_sql_err()?,
|
|
actual_total_cost_usd: row.try_get("actual_total_cost_usd").map_sql_err()?,
|
|
error_requests: row_u64(&row, "error_requests")?,
|
|
response_time_sum_ms: row.try_get("response_time_sum_ms").map_sql_err()?,
|
|
response_time_samples: row_u64(&row, "response_time_samples")?,
|
|
}))
|
|
}
|
|
|
|
pub async fn list_dashboard_daily_breakdown_from_daily_aggregates(
|
|
&self,
|
|
query: &UsageDashboardDailyBreakdownQuery,
|
|
) -> Result<Vec<StoredUsageDashboardDailyBreakdownRow>, DataLayerError> {
|
|
let rows = if let Some(user_id) = query.user_id.as_deref() {
|
|
sqlx::query(
|
|
r#"
|
|
SELECT
|
|
DATE_FORMAT(FROM_UNIXTIME(`date`), '%Y-%m-%d') AS date,
|
|
'aggregate' AS model,
|
|
'aggregate' AS provider,
|
|
CAST(COALESCE(SUM(total_requests), 0) AS SIGNED) AS requests,
|
|
CAST(COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) AS SIGNED) AS total_tokens,
|
|
CAST(COALESCE(SUM(COALESCE(total_cost, 0)), 0) AS DOUBLE) AS total_cost_usd,
|
|
CAST(0.0 AS DOUBLE) AS response_time_sum_ms,
|
|
CAST(0 AS SIGNED) AS response_time_samples
|
|
FROM stats_user_daily
|
|
WHERE user_id = ?
|
|
AND `date` >= ?
|
|
AND `date` < ?
|
|
AND total_requests > 0
|
|
GROUP BY `date`
|
|
ORDER BY `date` ASC
|
|
"#,
|
|
)
|
|
.bind(user_id)
|
|
.bind(to_i64(
|
|
query.created_from_unix_secs,
|
|
"stats_user_daily.date",
|
|
)?)
|
|
.bind(to_i64(
|
|
query.created_until_unix_secs,
|
|
"stats_user_daily.date",
|
|
)?)
|
|
.fetch_all(&self.pool)
|
|
.await
|
|
.map_sql_err()?
|
|
} else {
|
|
sqlx::query(
|
|
r#"
|
|
SELECT
|
|
DATE_FORMAT(FROM_UNIXTIME(`date`), '%Y-%m-%d') AS date,
|
|
'aggregate' AS model,
|
|
'aggregate' AS provider,
|
|
CAST(COALESCE(SUM(total_requests), 0) AS SIGNED) AS requests,
|
|
CAST(COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) AS SIGNED) AS total_tokens,
|
|
CAST(COALESCE(SUM(COALESCE(total_cost, 0)), 0) AS DOUBLE) AS total_cost_usd,
|
|
CAST(0.0 AS DOUBLE) AS response_time_sum_ms,
|
|
CAST(0 AS SIGNED) AS response_time_samples
|
|
FROM stats_daily
|
|
WHERE `date` >= ?
|
|
AND `date` < ?
|
|
AND total_requests > 0
|
|
GROUP BY `date`
|
|
ORDER BY `date` ASC
|
|
"#,
|
|
)
|
|
.bind(to_i64(query.created_from_unix_secs, "stats_daily.date")?)
|
|
.bind(to_i64(query.created_until_unix_secs, "stats_daily.date")?)
|
|
.fetch_all(&self.pool)
|
|
.await
|
|
.map_sql_err()?
|
|
};
|
|
|
|
rows.iter()
|
|
.map(|row| {
|
|
Ok(StoredUsageDashboardDailyBreakdownRow {
|
|
date: row.try_get("date").map_sql_err()?,
|
|
model: row.try_get("model").map_sql_err()?,
|
|
provider: row.try_get("provider").map_sql_err()?,
|
|
requests: row_u64(row, "requests")?,
|
|
total_tokens: row_u64(row, "total_tokens")?,
|
|
total_cost_usd: row.try_get("total_cost_usd").map_sql_err()?,
|
|
response_time_sum_ms: row.try_get("response_time_sum_ms").map_sql_err()?,
|
|
response_time_samples: row_u64(row, "response_time_samples")?,
|
|
})
|
|
})
|
|
.collect()
|
|
}
|
|
|
|
pub async fn summarize_usage_totals_by_user_ids(
|
|
&self,
|
|
user_ids: &[String],
|
|
) -> Result<Vec<StoredUsageUserTotals>, DataLayerError> {
|
|
if user_ids.is_empty() {
|
|
return Ok(Vec::new());
|
|
}
|
|
|
|
let unique_user_ids = user_ids
|
|
.iter()
|
|
.cloned()
|
|
.collect::<std::collections::BTreeSet<_>>();
|
|
let mut totals = BTreeMap::<String, StoredUsageUserTotals>::new();
|
|
let mut aggregate_cutoffs = BTreeMap::<String, u64>::new();
|
|
|
|
let mut aggregate_builder = QueryBuilder::<MySql>::new(
|
|
r#"
|
|
SELECT
|
|
user_id,
|
|
CAST(COALESCE(SUM(total_requests), 0) AS SIGNED) AS request_count,
|
|
CAST(COALESCE(
|
|
SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens),
|
|
0
|
|
) AS SIGNED) AS total_tokens,
|
|
MAX(`date`) AS latest_date
|
|
FROM stats_user_daily
|
|
WHERE user_id IN (
|
|
"#,
|
|
);
|
|
{
|
|
let mut separated = aggregate_builder.separated(", ");
|
|
for user_id in &unique_user_ids {
|
|
separated.push_bind(user_id.clone());
|
|
}
|
|
}
|
|
aggregate_builder.push(") GROUP BY user_id ORDER BY user_id ASC");
|
|
|
|
let aggregate_rows = aggregate_builder
|
|
.build()
|
|
.fetch_all(&self.pool)
|
|
.await
|
|
.map_sql_err()?;
|
|
for row in aggregate_rows {
|
|
let user_id: String = row.try_get("user_id").map_sql_err()?;
|
|
let latest_date = row.try_get::<i64, _>("latest_date").map_sql_err()?.max(0) as u64;
|
|
aggregate_cutoffs.insert(user_id.clone(), latest_date.saturating_add(86_400));
|
|
totals.insert(
|
|
user_id.clone(),
|
|
StoredUsageUserTotals {
|
|
user_id,
|
|
request_count: row_u64(&row, "request_count")?,
|
|
total_tokens: row_u64(&row, "total_tokens")?,
|
|
},
|
|
);
|
|
}
|
|
|
|
let mut raw_builder = QueryBuilder::<MySql>::new(format!(
|
|
r#"
|
|
SELECT
|
|
`usage`.user_id,
|
|
CAST(COUNT(*) AS SIGNED) AS request_count,
|
|
CAST(COALESCE(SUM({canonical_total_tokens_expr}), 0) AS SIGNED) AS total_tokens
|
|
FROM `usage`
|
|
JOIN (
|
|
"#,
|
|
canonical_total_tokens_expr = MYSQL_USAGE_CANONICAL_TOTAL_TOKENS_EXPR,
|
|
));
|
|
for (index, user_id) in unique_user_ids.iter().enumerate() {
|
|
if index > 0 {
|
|
raw_builder.push(" UNION ALL ");
|
|
}
|
|
let cutoff = aggregate_cutoffs.get(user_id).copied().unwrap_or_default();
|
|
raw_builder
|
|
.push("SELECT ")
|
|
.push_bind(user_id.clone())
|
|
.push(" AS user_id, ")
|
|
.push_bind(to_i64(cutoff, "usage aggregate cutoff")?)
|
|
.push(" AS cutoff_unix_secs");
|
|
}
|
|
raw_builder.push(
|
|
r#"
|
|
) AS requested ON requested.user_id = `usage`.user_id
|
|
LEFT JOIN usage_settlement_snapshots AS settlement
|
|
ON settlement.request_id = `usage`.request_id
|
|
WHERE `usage`.created_at_unix_ms >= requested.cutoff_unix_secs
|
|
AND `usage`.status NOT IN ('pending', 'streaming')
|
|
AND `usage`.provider_name NOT IN ('unknown', 'pending')
|
|
GROUP BY `usage`.user_id
|
|
ORDER BY `usage`.user_id ASC
|
|
"#,
|
|
);
|
|
|
|
let raw_rows = raw_builder
|
|
.build()
|
|
.fetch_all(&self.pool)
|
|
.await
|
|
.map_sql_err()?;
|
|
for row in raw_rows {
|
|
let user_id: String = row.try_get("user_id").map_sql_err()?;
|
|
let entry = totals
|
|
.entry(user_id.clone())
|
|
.or_insert_with(|| StoredUsageUserTotals {
|
|
user_id,
|
|
request_count: 0,
|
|
total_tokens: 0,
|
|
});
|
|
entry.request_count = entry
|
|
.request_count
|
|
.saturating_add(row_u64(&row, "request_count")?);
|
|
entry.total_tokens = entry
|
|
.total_tokens
|
|
.saturating_add(row_u64(&row, "total_tokens")?);
|
|
}
|
|
|
|
Ok(totals.into_values().collect())
|
|
}
|
|
}
|
|
|
|
impl MysqlUsageWriteRepository {
|
|
pub fn new(pool: MysqlPool) -> Self {
|
|
Self { pool }
|
|
}
|
|
|
|
pub async fn find_by_request_id(
|
|
&self,
|
|
request_id: &str,
|
|
) -> Result<Option<StoredRequestUsageAudit>, DataLayerError> {
|
|
let row = sqlx::query(&format!(
|
|
"{USAGE_COLUMNS} WHERE `usage`.request_id = ? LIMIT 1"
|
|
))
|
|
.bind(request_id)
|
|
.fetch_optional(&self.pool)
|
|
.await
|
|
.map_sql_err()?;
|
|
let usage = row
|
|
.as_ref()
|
|
.map(|row| map_usage_row(row, true))
|
|
.transpose()?;
|
|
match usage {
|
|
Some(usage) => http_capture::hydrate_usage_body_refs(&self.pool, usage)
|
|
.await
|
|
.map(Some),
|
|
None => Ok(None),
|
|
}
|
|
}
|
|
}
|
|
|
|
#[async_trait]
|
|
impl UsageWriteRepository for MysqlUsageWriteRepository {
|
|
async fn upsert(
|
|
&self,
|
|
usage: UpsertUsageRecord,
|
|
) -> Result<StoredRequestUsageAudit, DataLayerError> {
|
|
let mut usage = strip_deprecated_usage_display_fields(usage);
|
|
usage.validate()?;
|
|
let prepared_capture = http_capture::prepare_usage_http_capture(&mut usage)?;
|
|
let mut tx = self.pool.begin().await.map_sql_err()?;
|
|
let existing = counters::lock_and_load_usage(&mut tx, &usage.request_id).await?;
|
|
let recovers_terminal_failure = existing.as_ref().is_some_and(|existing| {
|
|
usage_can_recover_terminal_failure(
|
|
&existing.status,
|
|
&existing.billing_status,
|
|
&usage.status,
|
|
&usage.billing_status,
|
|
)
|
|
});
|
|
if let Some(existing) = existing.as_ref() {
|
|
if (existing.billing_status == "settled" || existing.billing_status == "void")
|
|
&& !recovers_terminal_failure
|
|
{
|
|
let existing = existing.clone();
|
|
tx.rollback().await.map_sql_err()?;
|
|
return http_capture::hydrate_usage_body_refs(&self.pool, existing).await;
|
|
}
|
|
}
|
|
|
|
let capture_update_allowed = recovers_terminal_failure
|
|
|| http_capture::capture_update_allowed(existing.as_ref(), &usage.status);
|
|
if capture_update_allowed {
|
|
http_capture::apply_previous_metadata_tombstones(&mut usage, existing.as_ref());
|
|
}
|
|
let prepared_snapshots = capture_update_allowed
|
|
.then(|| snapshots::from_usage(&usage))
|
|
.transpose()?;
|
|
bind_upsert(sqlx::query(UPSERT_USAGE_SQL), &usage)?
|
|
.execute(&mut *tx)
|
|
.await
|
|
.map_sql_err()?;
|
|
if capture_update_allowed {
|
|
http_capture::sync_usage_http_capture(&mut tx, &usage.request_id, &prepared_capture)
|
|
.await?;
|
|
let (routing_snapshot, settlement_snapshot) = prepared_snapshots
|
|
.as_ref()
|
|
.expect("capture-allowed usage has prepared snapshots");
|
|
snapshots::sync(
|
|
&mut tx,
|
|
&usage.request_id,
|
|
routing_snapshot,
|
|
settlement_snapshot,
|
|
matches!(usage.status.as_str(), "completed" | "failed" | "cancelled"),
|
|
)
|
|
.await?;
|
|
}
|
|
counters::enqueue_usage_transition_for_request(
|
|
&mut tx,
|
|
&usage.request_id,
|
|
existing.as_ref(),
|
|
)
|
|
.await?;
|
|
tx.commit().await.map_sql_err()?;
|
|
self.find_by_request_id(&usage.request_id)
|
|
.await?
|
|
.ok_or_else(|| {
|
|
DataLayerError::UnexpectedValue("usage upsert returned no row".to_string())
|
|
})
|
|
}
|
|
|
|
async fn rebuild_api_key_usage_stats(&self) -> Result<u64, DataLayerError> {
|
|
let mut tx = self.pool.begin().await.map_sql_err()?;
|
|
sqlx::query(
|
|
r#"
|
|
UPDATE api_keys
|
|
SET total_requests = 0,
|
|
total_tokens = 0,
|
|
total_cost_usd = 0,
|
|
last_used_at = NULL
|
|
"#,
|
|
)
|
|
.execute(&mut *tx)
|
|
.await
|
|
.map_sql_err()?;
|
|
|
|
let rows_affected = sqlx::query(&format!(
|
|
r#"
|
|
UPDATE api_keys
|
|
JOIN (
|
|
SELECT
|
|
api_key_id,
|
|
COUNT(*) AS total_requests,
|
|
COALESCE(SUM({canonical_total_tokens_expr}), 0) AS total_tokens,
|
|
COALESCE(SUM(COALESCE(total_cost_usd, 0)), 0) AS total_cost_usd,
|
|
MAX(created_at_unix_ms) AS last_used_at
|
|
FROM `usage`
|
|
LEFT JOIN usage_settlement_snapshots AS settlement
|
|
ON settlement.request_id = `usage`.request_id
|
|
WHERE api_key_id IS NOT NULL
|
|
AND TRIM(api_key_id) <> ''
|
|
AND status NOT IN ('pending', 'streaming')
|
|
GROUP BY api_key_id
|
|
) AS aggregated ON aggregated.api_key_id = api_keys.id
|
|
SET api_keys.total_requests = aggregated.total_requests,
|
|
api_keys.total_tokens = aggregated.total_tokens,
|
|
api_keys.total_cost_usd = aggregated.total_cost_usd,
|
|
api_keys.last_used_at = aggregated.last_used_at
|
|
"#,
|
|
canonical_total_tokens_expr = MYSQL_USAGE_CANONICAL_TOTAL_TOKENS_EXPR,
|
|
))
|
|
.execute(&mut *tx)
|
|
.await
|
|
.map_sql_err()?
|
|
.rows_affected();
|
|
tx.commit().await.map_sql_err()?;
|
|
Ok(rows_affected)
|
|
}
|
|
|
|
async fn rebuild_provider_api_key_usage_stats(&self) -> Result<u64, DataLayerError> {
|
|
let mut tx = self.pool.begin().await.map_sql_err()?;
|
|
sqlx::query(
|
|
r#"
|
|
UPDATE provider_api_keys
|
|
SET request_count = 0,
|
|
success_count = 0,
|
|
error_count = 0,
|
|
total_tokens = 0,
|
|
total_cost_usd = 0,
|
|
total_response_time_ms = 0,
|
|
last_used_at = NULL
|
|
"#,
|
|
)
|
|
.execute(&mut *tx)
|
|
.await
|
|
.map_sql_err()?;
|
|
|
|
let rows_affected = sqlx::query(&format!(
|
|
r#"
|
|
UPDATE provider_api_keys
|
|
JOIN (
|
|
SELECT
|
|
provider_api_key_id,
|
|
COUNT(*) AS request_count,
|
|
COALESCE(SUM({success_flag_expr}), 0) AS success_count,
|
|
COALESCE(SUM({error_flag_expr}), 0) AS error_count,
|
|
COALESCE(SUM(CASE
|
|
WHEN status IN ('pending', 'streaming') THEN 0
|
|
ELSE {canonical_total_tokens_expr}
|
|
END), 0) AS total_tokens,
|
|
COALESCE(SUM(CASE
|
|
WHEN status IN ('pending', 'streaming') THEN 0
|
|
ELSE COALESCE(total_cost_usd, 0)
|
|
END), 0) AS total_cost_usd,
|
|
COALESCE(SUM(CASE
|
|
WHEN {success_flag_expr} = 1 AND response_time_ms IS NOT NULL
|
|
THEN GREATEST(COALESCE(response_time_ms, 0), 0)
|
|
ELSE 0
|
|
END), 0) AS total_response_time_ms,
|
|
MAX(created_at_unix_ms) AS last_used_at
|
|
FROM `usage`
|
|
LEFT JOIN usage_settlement_snapshots AS settlement
|
|
ON settlement.request_id = `usage`.request_id
|
|
WHERE provider_api_key_id IS NOT NULL
|
|
AND TRIM(provider_api_key_id) <> ''
|
|
GROUP BY provider_api_key_id
|
|
) AS aggregated ON aggregated.provider_api_key_id = provider_api_keys.id
|
|
SET provider_api_keys.request_count = aggregated.request_count,
|
|
provider_api_keys.success_count = aggregated.success_count,
|
|
provider_api_keys.error_count = aggregated.error_count,
|
|
provider_api_keys.total_tokens = aggregated.total_tokens,
|
|
provider_api_keys.total_cost_usd = aggregated.total_cost_usd,
|
|
provider_api_keys.total_response_time_ms = aggregated.total_response_time_ms,
|
|
provider_api_keys.last_used_at = aggregated.last_used_at
|
|
"#,
|
|
success_flag_expr = MYSQL_PROVIDER_KEY_SUCCESS_FLAG_EXPR,
|
|
error_flag_expr = MYSQL_PROVIDER_KEY_ERROR_FLAG_EXPR,
|
|
canonical_total_tokens_expr = MYSQL_USAGE_CANONICAL_TOTAL_TOKENS_EXPR,
|
|
))
|
|
.execute(&mut *tx)
|
|
.await
|
|
.map_sql_err()?
|
|
.rows_affected();
|
|
tx.commit().await.map_sql_err()?;
|
|
Ok(rows_affected)
|
|
}
|
|
|
|
async fn cleanup_stale_pending_requests(
|
|
&self,
|
|
cutoff_unix_secs: u64,
|
|
now_unix_secs: u64,
|
|
timeout_minutes: u64,
|
|
batch_size: usize,
|
|
) -> Result<PendingUsageCleanupSummary, DataLayerError> {
|
|
if batch_size == 0 {
|
|
return Ok(PendingUsageCleanupSummary::default());
|
|
}
|
|
|
|
let now_unix_ms = now_unix_secs.saturating_mul(1000);
|
|
let mut summary = PendingUsageCleanupSummary::default();
|
|
let batch_size_u64 = u64::try_from(batch_size).map_err(|_| {
|
|
DataLayerError::InvalidInput(format!(
|
|
"invalid stale pending usage batch size: {batch_size}"
|
|
))
|
|
})?;
|
|
|
|
loop {
|
|
let mut tx = self.pool.begin().await.map_sql_err()?;
|
|
let stale_rows = sqlx::query(SELECT_STALE_PENDING_USAGE_BATCH_SQL)
|
|
.bind(to_i64(cutoff_unix_secs, "stale pending usage cutoff")?)
|
|
.bind(to_i64(batch_size_u64, "stale pending usage batch size")?)
|
|
.fetch_all(&mut *tx)
|
|
.await
|
|
.map_sql_err()?;
|
|
|
|
if stale_rows.is_empty() {
|
|
tx.rollback().await.map_sql_err()?;
|
|
break;
|
|
}
|
|
|
|
let stale_rows = stale_rows
|
|
.iter()
|
|
.map(|row| {
|
|
Ok(StalePendingUsageRow {
|
|
request_id: row.try_get("request_id").map_sql_err()?,
|
|
status: row.try_get("status").map_sql_err()?,
|
|
billing_status: row.try_get("billing_status").map_sql_err()?,
|
|
})
|
|
})
|
|
.collect::<Result<Vec<_>, DataLayerError>>()?;
|
|
let completed_request_ids =
|
|
completed_request_ids_mysql(&mut tx, stale_rows.iter().map(|row| &row.request_id))
|
|
.await?;
|
|
|
|
for row in stale_rows {
|
|
if completed_request_ids.contains(&row.request_id) {
|
|
sqlx::query(
|
|
r#"
|
|
UPDATE `usage`
|
|
SET status = 'completed',
|
|
status_code = 200,
|
|
error_message = NULL
|
|
WHERE request_id = ?
|
|
"#,
|
|
)
|
|
.bind(&row.request_id)
|
|
.execute(&mut *tx)
|
|
.await
|
|
.map_sql_err()?;
|
|
sqlx::query(
|
|
r#"
|
|
UPDATE request_candidates
|
|
SET status = 'success',
|
|
finished_at = ?
|
|
WHERE request_id = ?
|
|
AND status = 'streaming'
|
|
"#,
|
|
)
|
|
.bind(to_i64(now_unix_ms, "request candidate finished_at")?)
|
|
.bind(&row.request_id)
|
|
.execute(&mut *tx)
|
|
.await
|
|
.map_sql_err()?;
|
|
summary.recovered += 1;
|
|
continue;
|
|
}
|
|
|
|
let candidate_info =
|
|
latest_failed_candidate_mysql(&mut tx, &row.request_id).await?;
|
|
let (status_code, error_message) = resolve_stale_pending_failure(
|
|
candidate_info.as_ref(),
|
|
&row.status,
|
|
timeout_minutes,
|
|
);
|
|
let status_code_i64 = i64::from(status_code);
|
|
if row.billing_status == "pending" {
|
|
sqlx::query(
|
|
r#"
|
|
UPDATE `usage`
|
|
SET status = 'failed',
|
|
status_code = ?,
|
|
error_message = ?,
|
|
billing_status = 'void',
|
|
finalized_at = ?,
|
|
total_cost_usd = 0,
|
|
actual_total_cost_usd = 0
|
|
WHERE request_id = ?
|
|
"#,
|
|
)
|
|
.bind(status_code_i64)
|
|
.bind(&error_message)
|
|
.bind(to_i64(now_unix_secs, "usage finalized_at")?)
|
|
.bind(&row.request_id)
|
|
.execute(&mut *tx)
|
|
.await
|
|
.map_sql_err()?;
|
|
upsert_void_usage_settlement_snapshot_mysql(
|
|
&mut tx,
|
|
&row.request_id,
|
|
now_unix_secs,
|
|
)
|
|
.await?;
|
|
} else {
|
|
sqlx::query(
|
|
r#"
|
|
UPDATE `usage`
|
|
SET status = 'failed',
|
|
status_code = ?,
|
|
error_message = ?
|
|
WHERE request_id = ?
|
|
"#,
|
|
)
|
|
.bind(status_code_i64)
|
|
.bind(&error_message)
|
|
.bind(&row.request_id)
|
|
.execute(&mut *tx)
|
|
.await
|
|
.map_sql_err()?;
|
|
}
|
|
|
|
sqlx::query(
|
|
r#"
|
|
UPDATE request_candidates
|
|
SET status = 'failed',
|
|
finished_at = ?,
|
|
error_message = '请求超时（服务器可能已重启）'
|
|
WHERE request_id = ?
|
|
AND status IN ('pending', 'streaming')
|
|
"#,
|
|
)
|
|
.bind(to_i64(now_unix_ms, "request candidate finished_at")?)
|
|
.bind(&row.request_id)
|
|
.execute(&mut *tx)
|
|
.await
|
|
.map_sql_err()?;
|
|
summary.failed += 1;
|
|
}
|
|
|
|
tx.commit().await.map_sql_err()?;
|
|
}
|
|
|
|
Ok(summary)
|
|
}
|
|
|
|
async fn flush_usage_counter_deltas(
|
|
&self,
|
|
batch_size: usize,
|
|
) -> Result<aether_data_contracts::repository::usage::UsageCounterFlushSummary, DataLayerError>
|
|
{
|
|
counters::flush(&self.pool, batch_size).await
|
|
}
|
|
|
|
async fn enqueue_proxy_node_counter_delta(
|
|
&self,
|
|
delta: aether_data_contracts::repository::usage::ProxyNodeCounterDelta,
|
|
) -> Result<bool, DataLayerError> {
|
|
counters::enqueue_proxy_node(&self.pool, delta).await
|
|
}
|
|
|
|
async fn enqueue_management_token_counter_delta(
|
|
&self,
|
|
delta: aether_data_contracts::repository::usage::ManagementTokenCounterDelta,
|
|
) -> Result<bool, DataLayerError> {
|
|
counters::enqueue_management_token(&self.pool, delta).await
|
|
}
|
|
|
|
async fn enqueue_api_key_last_used_delta(
|
|
&self,
|
|
delta: aether_data_contracts::repository::usage::ApiKeyLastUsedDelta,
|
|
) -> Result<bool, DataLayerError> {
|
|
counters::enqueue_api_key_last_used(&self.pool, delta).await
|
|
}
|
|
|
|
async fn cleanup_processed_usage_counter_deltas(
|
|
&self,
|
|
cutoff_unix_secs: u64,
|
|
batch_size: usize,
|
|
) -> Result<usize, DataLayerError> {
|
|
counters::cleanup_processed(&self.pool, cutoff_unix_secs, batch_size).await
|
|
}
|
|
|
|
async fn cleanup_usage(
|
|
&self,
|
|
window: &UsageCleanupWindow,
|
|
batch_size: usize,
|
|
auto_delete_expired_keys: bool,
|
|
targets: UsageCleanupTargets,
|
|
mode: UsageCleanupExecutionMode,
|
|
) -> Result<UsageCleanupSummary, DataLayerError> {
|
|
cleanup::cleanup_usage(
|
|
&self.pool,
|
|
window,
|
|
batch_size,
|
|
auto_delete_expired_keys,
|
|
targets,
|
|
mode,
|
|
)
|
|
.await
|
|
}
|
|
|
|
async fn preview_usage_cleanup(
|
|
&self,
|
|
window: &UsageCleanupWindow,
|
|
targets: UsageCleanupTargets,
|
|
mode: UsageCleanupExecutionMode,
|
|
) -> Result<UsageCleanupPreviewCounts, DataLayerError> {
|
|
cleanup::preview_usage_cleanup(&self.pool, window, targets, mode).await
|
|
}
|
|
}
|
|
|
|
struct StalePendingUsageRow {
|
|
request_id: String,
|
|
status: String,
|
|
billing_status: String,
|
|
}
|
|
|
|
async fn completed_request_ids_mysql<'a>(
|
|
tx: &mut sqlx::Transaction<'_, sqlx::MySql>,
|
|
request_ids: impl Iterator<Item = &'a String>,
|
|
) -> Result<HashSet<String>, DataLayerError> {
|
|
let mut completed = HashSet::new();
|
|
for request_id in request_ids {
|
|
let rows = sqlx::query(SELECT_COMPLETED_REQUEST_CANDIDATES_SQL)
|
|
.bind(request_id)
|
|
.fetch_all(&mut **tx)
|
|
.await
|
|
.map_sql_err()?;
|
|
let mut is_completed = false;
|
|
for row in &rows {
|
|
if candidate_row_is_completed(row)? {
|
|
is_completed = true;
|
|
break;
|
|
}
|
|
}
|
|
if is_completed {
|
|
completed.insert(request_id.clone());
|
|
}
|
|
}
|
|
Ok(completed)
|
|
}
|
|
|
|
fn candidate_row_is_completed(row: &MySqlRow) -> Result<bool, DataLayerError> {
|
|
let status: String = row.try_get("status").map_sql_err()?;
|
|
if status == "streaming" {
|
|
return Ok(true);
|
|
}
|
|
if status != "success" {
|
|
return Ok(false);
|
|
}
|
|
let Some(extra_data) = row
|
|
.try_get::<Option<String>, _>("extra_data")
|
|
.map_sql_err()?
|
|
else {
|
|
return Ok(false);
|
|
};
|
|
let Ok(value) = serde_json::from_str::<serde_json::Value>(&extra_data) else {
|
|
return Ok(false);
|
|
};
|
|
Ok(value
|
|
.get("stream_completed")
|
|
.and_then(serde_json::Value::as_bool)
|
|
.unwrap_or(false))
|
|
}
|
|
|
|
async fn upsert_void_usage_settlement_snapshot_mysql(
|
|
tx: &mut sqlx::Transaction<'_, sqlx::MySql>,
|
|
request_id: &str,
|
|
now_unix_secs: u64,
|
|
) -> Result<(), DataLayerError> {
|
|
let now = to_i64(now_unix_secs, "usage settlement snapshot timestamp")?;
|
|
sqlx::query(
|
|
r#"
|
|
INSERT INTO usage_settlement_snapshots (
|
|
request_id,
|
|
billing_status,
|
|
finalized_at,
|
|
created_at,
|
|
updated_at
|
|
) VALUES (?, 'void', ?, ?, ?)
|
|
ON DUPLICATE KEY UPDATE
|
|
billing_status = VALUES(billing_status),
|
|
billing_snapshot_schema_version = NULL,
|
|
billing_snapshot_status = NULL,
|
|
settlement_snapshot_schema_version = NULL,
|
|
settlement_snapshot = NULL,
|
|
billing_dimensions = NULL,
|
|
billing_input_tokens = NULL,
|
|
billing_effective_input_tokens = NULL,
|
|
billing_output_tokens = NULL,
|
|
billing_cache_creation_tokens = NULL,
|
|
billing_cache_creation_5m_tokens = NULL,
|
|
billing_cache_creation_1h_tokens = NULL,
|
|
billing_cache_read_tokens = NULL,
|
|
billing_total_input_context = NULL,
|
|
billing_cache_creation_cost_usd = NULL,
|
|
billing_cache_read_cost_usd = NULL,
|
|
billing_total_cost_usd = NULL,
|
|
billing_actual_total_cost_usd = NULL,
|
|
billing_pricing_source = NULL,
|
|
billing_rule_id = NULL,
|
|
billing_rule_version = NULL,
|
|
rate_multiplier = NULL,
|
|
is_free_tier = NULL,
|
|
input_price_per_1m = NULL,
|
|
output_price_per_1m = NULL,
|
|
cache_creation_price_per_1m = NULL,
|
|
cache_read_price_per_1m = NULL,
|
|
price_per_request = NULL,
|
|
finalized_at = COALESCE(usage_settlement_snapshots.finalized_at, VALUES(finalized_at)),
|
|
updated_at = VALUES(updated_at)
|
|
"#,
|
|
)
|
|
.bind(request_id)
|
|
.bind(now)
|
|
.bind(now)
|
|
.bind(now)
|
|
.execute(&mut **tx)
|
|
.await
|
|
.map_sql_err()?;
|
|
Ok(())
|
|
}
|
|
|
|
fn stale_pending_error_message(status: &str, timeout_minutes: u64) -> String {
|
|
format!("请求超时: 状态 '{status}' 超过 {timeout_minutes} 分钟未完成")
|
|
}
|
|
|
|
struct FailedCandidateCleanupInfo {
|
|
status_code: Option<u16>,
|
|
error_message: Option<String>,
|
|
}
|
|
|
|
fn resolve_stale_pending_failure(
|
|
candidate: Option<&FailedCandidateCleanupInfo>,
|
|
status: &str,
|
|
timeout_minutes: u64,
|
|
) -> (u16, String) {
|
|
match candidate {
|
|
Some(info) => (
|
|
info.status_code.unwrap_or(502),
|
|
info.error_message
|
|
.clone()
|
|
.unwrap_or_else(|| stale_pending_error_message(status, timeout_minutes)),
|
|
),
|
|
None => (504, stale_pending_error_message(status, timeout_minutes)),
|
|
}
|
|
}
|
|
|
|
async fn latest_failed_candidate_mysql(
|
|
tx: &mut sqlx::Transaction<'_, sqlx::MySql>,
|
|
request_id: &str,
|
|
) -> Result<Option<FailedCandidateCleanupInfo>, DataLayerError> {
|
|
let row = sqlx::query(
|
|
r#"
|
|
SELECT status_code, error_message
|
|
FROM request_candidates
|
|
WHERE request_id = ?
|
|
AND status IN ('failed', 'cancelled')
|
|
ORDER BY
|
|
COALESCE(finished_at, started_at, created_at) DESC,
|
|
retry_index DESC,
|
|
candidate_index DESC
|
|
LIMIT 1
|
|
"#,
|
|
)
|
|
.bind(request_id)
|
|
.fetch_optional(&mut **tx)
|
|
.await
|
|
.map_sql_err()?;
|
|
|
|
let Some(row) = row else {
|
|
return Ok(None);
|
|
};
|
|
let status_code = row
|
|
.try_get::<Option<i64>, _>("status_code")
|
|
.map_sql_err()?
|
|
.and_then(|value| u16::try_from(value).ok());
|
|
let error_message = row
|
|
.try_get::<Option<String>, _>("error_message")
|
|
.map_sql_err()?
|
|
.map(|value| value.trim().to_string())
|
|
.filter(|value| !value.is_empty());
|
|
Ok(Some(FailedCandidateCleanupInfo {
|
|
status_code,
|
|
error_message,
|
|
}))
|
|
}
|
|
|
|
fn bind_upsert<'q>(
|
|
mut query: sqlx::query::Query<'q, sqlx::MySql, sqlx::mysql::MySqlArguments>,
|
|
usage: &'q UpsertUsageRecord,
|
|
) -> Result<sqlx::query::Query<'q, sqlx::MySql, sqlx::mysql::MySqlArguments>, DataLayerError> {
|
|
let input_tokens = usage.input_tokens.unwrap_or_default();
|
|
let output_tokens = usage.output_tokens.unwrap_or_default();
|
|
let cache_creation_tokens = usage
|
|
.cache_creation_input_tokens
|
|
.or_else(|| {
|
|
Some(
|
|
usage
|
|
.cache_creation_ephemeral_5m_input_tokens
|
|
.unwrap_or_default()
|
|
+ usage
|
|
.cache_creation_ephemeral_1h_input_tokens
|
|
.unwrap_or_default(),
|
|
)
|
|
})
|
|
.unwrap_or_default();
|
|
let cache_read_tokens = usage.cache_read_input_tokens.unwrap_or_default();
|
|
let total_tokens = usage
|
|
.total_tokens
|
|
.unwrap_or(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens);
|
|
let created_at = usage
|
|
.created_at_unix_ms
|
|
.unwrap_or(usage.updated_at_unix_secs);
|
|
let request_metadata = usage
|
|
.request_metadata
|
|
.as_ref()
|
|
.map(serde_json::to_string)
|
|
.transpose()
|
|
.map_err(|err| DataLayerError::InvalidInput(err.to_string()))?;
|
|
|
|
query = query
|
|
.bind(&usage.request_id)
|
|
.bind(&usage.request_id)
|
|
.bind(usage.user_id.as_deref())
|
|
.bind(usage.api_key_id.as_deref())
|
|
.bind(&usage.provider_name)
|
|
.bind(&usage.model)
|
|
.bind(usage.target_model.as_deref())
|
|
.bind(usage.provider_id.as_deref())
|
|
.bind(usage.provider_endpoint_id.as_deref())
|
|
.bind(usage.provider_api_key_id.as_deref())
|
|
.bind(usage.request_type.as_deref())
|
|
.bind(usage.api_format.as_deref())
|
|
.bind(usage.api_family.as_deref())
|
|
.bind(usage.endpoint_kind.as_deref())
|
|
.bind(usage.endpoint_api_format.as_deref())
|
|
.bind(usage.provider_api_family.as_deref())
|
|
.bind(usage.provider_endpoint_kind.as_deref())
|
|
.bind(usage.has_format_conversion.unwrap_or(false))
|
|
.bind(usage.is_stream.unwrap_or(false))
|
|
.bind(usage_upstream_is_stream(usage))
|
|
.bind(to_i64(input_tokens, "input_tokens")?)
|
|
.bind(to_i64(output_tokens, "output_tokens")?)
|
|
.bind(to_i64(total_tokens, "total_tokens")?)
|
|
.bind(to_i64(
|
|
cache_creation_tokens,
|
|
"cache_creation_input_tokens",
|
|
)?)
|
|
.bind(to_i64(
|
|
usage
|
|
.cache_creation_ephemeral_5m_input_tokens
|
|
.unwrap_or_default(),
|
|
"cache_creation_ephemeral_5m_input_tokens",
|
|
)?)
|
|
.bind(to_i64(
|
|
usage
|
|
.cache_creation_ephemeral_1h_input_tokens
|
|
.unwrap_or_default(),
|
|
"cache_creation_ephemeral_1h_input_tokens",
|
|
)?)
|
|
.bind(to_i64(cache_read_tokens, "cache_read_input_tokens")?)
|
|
.bind(usage.cache_creation_cost_usd.unwrap_or_default())
|
|
.bind(usage.cache_read_cost_usd.unwrap_or_default())
|
|
.bind(usage.output_price_per_1m)
|
|
.bind(usage.total_cost_usd.unwrap_or_default())
|
|
.bind(usage.actual_total_cost_usd.unwrap_or_default())
|
|
.bind(usage.status_code.map(i64::from))
|
|
.bind(usage.error_message.as_deref())
|
|
.bind(usage.error_category.as_deref())
|
|
.bind(usage.response_time_ms.map(|value| value as i64))
|
|
.bind(usage.first_byte_time_ms.map(|value| value as i64))
|
|
.bind(&usage.status)
|
|
.bind(&usage.billing_status)
|
|
.bind(request_metadata)
|
|
.bind(usage.candidate_id.as_deref())
|
|
.bind(usage.candidate_index.map(|value| value as i64))
|
|
.bind(usage.key_name.as_deref())
|
|
.bind(usage.planner_kind.as_deref())
|
|
.bind(usage.route_family.as_deref())
|
|
.bind(usage.route_kind.as_deref())
|
|
.bind(usage.execution_path.as_deref())
|
|
.bind(usage.local_execution_runtime_miss_reason.as_deref())
|
|
.bind(usage.finalized_at_unix_secs.map(|value| value as i64))
|
|
.bind(to_i64(created_at, "created_at_unix_ms")?)
|
|
.bind(to_i64(usage.updated_at_unix_secs, "updated_at_unix_secs")?);
|
|
Ok(query)
|
|
}
|
|
|
|
fn map_usage_row(
|
|
row: &MySqlRow,
|
|
resolve_legacy_compressed: bool,
|
|
) -> Result<StoredRequestUsageAudit, DataLayerError> {
|
|
let id = row
|
|
.try_get::<Option<String>, _>("id")
|
|
.map_sql_err()?
|
|
.unwrap_or_else(|| {
|
|
row.try_get::<String, _>("request_id")
|
|
.unwrap_or_else(|_| "unknown".to_string())
|
|
});
|
|
let mut audit = StoredRequestUsageAudit::new(
|
|
id,
|
|
row.try_get("request_id").map_sql_err()?,
|
|
row.try_get("user_id").map_sql_err()?,
|
|
row.try_get("api_key_id").map_sql_err()?,
|
|
row.try_get("username").map_sql_err()?,
|
|
row.try_get("api_key_name").map_sql_err()?,
|
|
row.try_get("provider_name").map_sql_err()?,
|
|
row.try_get("model").map_sql_err()?,
|
|
row.try_get("target_model").map_sql_err()?,
|
|
row.try_get("provider_id").map_sql_err()?,
|
|
row.try_get("provider_endpoint_id").map_sql_err()?,
|
|
row.try_get("provider_api_key_id").map_sql_err()?,
|
|
row.try_get("request_type").map_sql_err()?,
|
|
row.try_get("api_format").map_sql_err()?,
|
|
row.try_get("api_family").map_sql_err()?,
|
|
row.try_get("endpoint_kind").map_sql_err()?,
|
|
row.try_get("endpoint_api_format").map_sql_err()?,
|
|
row.try_get("provider_api_family").map_sql_err()?,
|
|
row.try_get("provider_endpoint_kind").map_sql_err()?,
|
|
row.try_get::<bool, _>("has_format_conversion")
|
|
.map_sql_err()?,
|
|
row.try_get::<bool, _>("is_stream").map_sql_err()?,
|
|
row_i32(row, "input_tokens")?,
|
|
row_i32(row, "output_tokens")?,
|
|
row_i32(row, "total_tokens")?,
|
|
row.try_get("total_cost_usd").map_sql_err()?,
|
|
row.try_get("actual_total_cost_usd").map_sql_err()?,
|
|
row_optional_i32(row, "status_code")?,
|
|
row.try_get("error_message").map_sql_err()?,
|
|
row.try_get("error_category").map_sql_err()?,
|
|
row_optional_i32(row, "response_time_ms")?,
|
|
row_optional_i32(row, "first_byte_time_ms")?,
|
|
row.try_get("status").map_sql_err()?,
|
|
row.try_get("billing_status").map_sql_err()?,
|
|
row.try_get("created_at_unix_ms").map_sql_err()?,
|
|
row.try_get("updated_at_unix_secs").map_sql_err()?,
|
|
row.try_get("finalized_at_unix_secs").map_sql_err()?,
|
|
)?;
|
|
audit.cache_creation_input_tokens = row_u64(row, "cache_creation_input_tokens")?;
|
|
audit.cache_creation_ephemeral_5m_input_tokens =
|
|
row_u64(row, "cache_creation_ephemeral_5m_input_tokens")?;
|
|
audit.cache_creation_ephemeral_1h_input_tokens =
|
|
row_u64(row, "cache_creation_ephemeral_1h_input_tokens")?;
|
|
audit.cache_read_input_tokens = row_u64(row, "cache_read_input_tokens")?;
|
|
audit.cache_creation_cost_usd = row.try_get("cache_creation_cost_usd").map_sql_err()?;
|
|
audit.cache_read_cost_usd = row.try_get("cache_read_cost_usd").map_sql_err()?;
|
|
audit.output_price_per_1m = row.try_get("output_price_per_1m").map_sql_err()?;
|
|
audit.request_metadata = row
|
|
.try_get::<Option<String>, _>("request_metadata")
|
|
.map_sql_err()?
|
|
.map(|raw| serde_json::from_str(&raw))
|
|
.transpose()
|
|
.map_err(|err| DataLayerError::UnexpectedValue(err.to_string()))?;
|
|
audit.client_family = usage_request_metadata_client_family(audit.request_metadata.as_ref())
|
|
.map(ToOwned::to_owned);
|
|
http_capture::hydrate_usage_row(row, &mut audit, resolve_legacy_compressed)?;
|
|
let upstream_is_stream = row
|
|
.try_get::<Option<bool>, _>("upstream_is_stream")
|
|
.map_sql_err()?;
|
|
merge_usage_stream_metadata(&mut audit.request_metadata, upstream_is_stream);
|
|
snapshots::hydrate_row(row, &mut audit)?;
|
|
Ok(audit)
|
|
}
|
|
|
|
fn to_i64(value: u64, field: &str) -> Result<i64, DataLayerError> {
|
|
i64::try_from(value).map_err(|_| DataLayerError::InvalidInput(format!("{field} overflow")))
|
|
}
|
|
|
|
fn usage_upstream_is_stream(usage: &UpsertUsageRecord) -> bool {
|
|
usage
|
|
.request_metadata
|
|
.as_ref()
|
|
.and_then(serde_json::Value::as_object)
|
|
.and_then(|metadata| metadata.get(UPSTREAM_IS_STREAM_KEY))
|
|
.and_then(serde_json::Value::as_bool)
|
|
.unwrap_or_else(|| usage.is_stream.unwrap_or(false))
|
|
}
|
|
|
|
fn merge_usage_stream_metadata(metadata: &mut Option<serde_json::Value>, upstream: Option<bool>) {
|
|
let Some(upstream) = upstream else {
|
|
return;
|
|
};
|
|
let value = metadata.get_or_insert_with(|| serde_json::json!({}));
|
|
let Some(object) = value.as_object_mut() else {
|
|
return;
|
|
};
|
|
object
|
|
.entry(UPSTREAM_IS_STREAM_KEY)
|
|
.or_insert(serde_json::Value::Bool(upstream));
|
|
}
|
|
|
|
fn row_i32(row: &MySqlRow, field: &str) -> Result<i32, DataLayerError> {
|
|
let value: i64 = row.try_get(field).map_sql_err()?;
|
|
i32::try_from(value).map_err(|_| DataLayerError::UnexpectedValue(format!("{field} overflow")))
|
|
}
|
|
|
|
fn row_optional_i32(row: &MySqlRow, field: &str) -> Result<Option<i32>, DataLayerError> {
|
|
row.try_get::<Option<i64>, _>(field)
|
|
.map_sql_err()?
|
|
.map(|value| {
|
|
i32::try_from(value)
|
|
.map_err(|_| DataLayerError::UnexpectedValue(format!("{field} overflow")))
|
|
})
|
|
.transpose()
|
|
}
|
|
|
|
fn row_u64(row: &MySqlRow, field: &str) -> Result<u64, DataLayerError> {
|
|
let value: i64 = row.try_get(field).map_sql_err()?;
|
|
u64::try_from(value).map_err(|_| DataLayerError::UnexpectedValue(format!("{field} negative")))
|
|
}
|
|
|
|
fn map_mysql_usage_daily_summary(
|
|
row: &MySqlRow,
|
|
) -> Result<StoredUsageDailySummary, DataLayerError> {
|
|
Ok(StoredUsageDailySummary {
|
|
date: row.try_get("date").map_sql_err()?,
|
|
requests: row_u64(row, "requests")?,
|
|
total_tokens: row_u64(row, "total_tokens")?,
|
|
total_cost_usd: row.try_get("total_cost_usd").map_sql_err()?,
|
|
actual_total_cost_usd: row.try_get("actual_total_cost_usd").map_sql_err()?,
|
|
})
|
|
}
|
|
|
|
fn usage_current_unix_secs() -> u64 {
|
|
SystemTime::now()
|
|
.duration_since(UNIX_EPOCH)
|
|
.map(|duration| duration.as_secs())
|
|
.unwrap_or_default()
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests;
|