mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-08 18:37:46 +08:00
5476 lines
212 KiB
Rust
5476 lines
212 KiB
Rust
use std::collections::BTreeMap;
|
|||
use std::sync::Arc;
|
|||
use std::sync::RwLock;
|
|||
|
|
|
||
use aether_ai_formats::UPSTREAM_IS_STREAM_KEY;
|
|||
use aether_data_contracts::repository::usage::{
|
|||
parse_usage_body_ref, usage_body_ref, StoredUsageAuditAggregation, StoredUsageAuditSummary,
|
|||
StoredUsageBreakdownSummaryRow, StoredUsageCacheAffinityHitSummary,
|
|||
|
|
StoredUsageCacheAffinityIntervalRow, StoredUsageCacheHitSummary, StoredUsageCostSavingsSummary,
|
||
|
|
StoredUsageDashboardDailyBreakdownRow, StoredUsageDashboardProviderCount,
|
||
|
|
StoredUsageDashboardSummary, StoredUsageErrorDistributionRow, StoredUsageLeaderboardSummary,
|
||
StoredUsagePerformancePercentilesRow, StoredUsageProviderPerformance,
|
|||
|
|
StoredUsageProviderPerformanceProviderRow, StoredUsageProviderPerformanceSummary,
|
||
|
|
StoredUsageProviderPerformanceTimelineRow, StoredUsageSettledCostSummary,
|
||
StoredUsageTimeSeriesBucket, StoredUsageUserTotals, UsageAuditAggregationGroupBy,
|
|||
|
|
UsageAuditAggregationQuery, UsageAuditKeywordSearchQuery, UsageAuditSummaryQuery,
|
||
|
|
UsageBodyField, UsageBreakdownGroupBy, UsageBreakdownSummaryQuery,
|
||
|
|
UsageCacheAffinityHitSummaryQuery, UsageCacheAffinityIntervalGroupBy,
|
||
|
|
UsageCacheAffinityIntervalQuery, UsageCacheHitSummaryQuery, UsageCostSavingsSummaryQuery,
|
||
|
|
UsageDashboardDailyBreakdownQuery, UsageDashboardProviderCountsQuery,
|
||
|
|
UsageDashboardSummaryQuery, UsageErrorDistributionQuery, UsageLeaderboardGroupBy,
|
||
|
|
UsageLeaderboardQuery, UsageMonitoringErrorCountQuery, UsageMonitoringErrorListQuery,
|
||
UsagePerformancePercentilesQuery, UsageProviderPerformanceQuery, UsageSettledCostSummaryQuery,
|
|||
|
|
UsageTimeSeriesGranularity, UsageTimeSeriesQuery,
|
||
};
|
|||
use async_trait::async_trait;
|
|||
use chrono::Utc;
|
|||
use serde_json::Value;
|
|||
|
|||
use super::{
|
|||
api_key_usage_contribution, provider_api_key_usage_contribution,
|
|||
|
|
strip_deprecated_usage_display_fields, usage_can_recover_terminal_failure,
|
||
usage_request_metadata_client_family, ApiKeyUsageContribution, ApiKeyUsageDelta,
|
|||
|
|
ProviderApiKeyUsageContribution, ProviderApiKeyUsageDelta, ProviderApiKeyWindowUsageRequest,
|
||
|
|
StoredProviderApiKeyUsageSummary, StoredProviderApiKeyWindowUsageSummary,
|
||
|
|
StoredProviderUsageSummary, StoredProviderUsageWindow, StoredRequestUsageAudit,
|
||
|
|
StoredUsageDailySummary, UpsertUsageRecord, UsageAuditListQuery, UsageDailyHeatmapQuery,
|
||
|
|
UsageReadRepository, UsageWriteRepository,
|
||
};
|
|||
use crate::repository::auth::InMemoryAuthApiKeySnapshotRepository;
|
|||
use crate::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
|
|||
use crate::DataLayerError;
|
|||
|
|
|
||
|
|
#[derive(Debug, Default)]
|
||
|
|
pub struct InMemoryUsageReadRepository {
|
||
|
|
by_request_id: RwLock<BTreeMap<String, StoredRequestUsageAudit>>,
|
||
detached_bodies: RwLock<BTreeMap<String, Value>>,
|
|||
provider_usage_windows: RwLock<Vec<StoredProviderUsageWindow>>,
|
|||
auth_api_keys: Option<Arc<InMemoryAuthApiKeySnapshotRepository>>,
|
|||
provider_catalog: Option<Arc<InMemoryProviderCatalogReadRepository>>,
|
|||
}
|
|||
|
|
|
||
|
|
impl InMemoryUsageReadRepository {
|
||
|
|
pub fn seed<I>(items: I) -> Self
|
||
|
|
where
|
||
|
|
I: IntoIterator<Item = StoredRequestUsageAudit>,
|
||
|
|
{
|
||
|
|
let mut by_request_id = BTreeMap::new();
|
||
for mut item in items {
|
|||
|
|
hydrate_legacy_body_refs(&mut item);
|
||
hydrate_client_family(&mut item);
|
|||
by_request_id.insert(item.request_id.clone(), item);
|
|||
|
|
}
|
||
|
|
Self {
|
||
|
|
by_request_id: RwLock::new(by_request_id),
|
||
detached_bodies: RwLock::new(BTreeMap::new()),
|
|||
|
|
provider_usage_windows: RwLock::new(Vec::new()),
|
||
auth_api_keys: None,
|
|||
provider_catalog: None,
|
|||
}
|
|||
|
|
}
|
||
|
|
|
||
|
|
pub fn seed_with_detached_bodies<I>(items: I) -> Self
|
||
|
|
where
|
||
|
|
I: IntoIterator<Item = StoredRequestUsageAudit>,
|
||
|
|
{
|
||
|
|
let mut by_request_id = BTreeMap::new();
|
||
|
|
let mut detached_bodies = BTreeMap::new();
|
||
|
|
for mut item in items {
|
||
|
|
hydrate_legacy_body_refs(&mut item);
|
||
hydrate_client_family(&mut item);
|
|||
let request_id = item.request_id.clone();
|
|||
|
|
if let Some(body_ref) = detach_usage_body(
|
||
|
|
&request_id,
|
||
|
|
&mut item.request_body,
|
||
|
|
&mut detached_bodies,
|
||
|
|
UsageBodyField::RequestBody,
|
||
|
|
) {
|
||
|
|
item.request_body_ref = Some(body_ref);
|
||
|
|
}
|
||
|
|
if let Some(body_ref) = detach_usage_body(
|
||
|
|
&request_id,
|
||
|
|
&mut item.provider_request_body,
|
||
|
|
&mut detached_bodies,
|
||
|
|
UsageBodyField::ProviderRequestBody,
|
||
|
|
) {
|
||
|
|
item.provider_request_body_ref = Some(body_ref);
|
||
|
|
}
|
||
|
|
if let Some(body_ref) = detach_usage_body(
|
||
|
|
&request_id,
|
||
|
|
&mut item.response_body,
|
||
|
|
&mut detached_bodies,
|
||
|
|
UsageBodyField::ResponseBody,
|
||
|
|
) {
|
||
|
|
item.response_body_ref = Some(body_ref);
|
||
|
|
}
|
||
|
|
if let Some(body_ref) = detach_usage_body(
|
||
|
|
&request_id,
|
||
|
|
&mut item.client_response_body,
|
||
|
|
&mut detached_bodies,
|
||
|
|
UsageBodyField::ClientResponseBody,
|
||
|
|
) {
|
||
|
|
item.client_response_body_ref = Some(body_ref);
|
||
|
|
}
|
||
|
|
by_request_id.insert(request_id, item);
|
||
|
|
}
|
||
|
|
Self {
|
||
|
|
by_request_id: RwLock::new(by_request_id),
|
||
|
|
detached_bodies: RwLock::new(detached_bodies),
|
||
provider_usage_windows: RwLock::new(Vec::new()),
|
|||
auth_api_keys: None,
|
|||
provider_catalog: None,
|
|||
}
|
|||
|
|
}
|
||
|
|
|
||
|
|
pub fn with_provider_usage_windows<I>(self, items: I) -> Self
|
||
|
|
where
|
||
|
|
I: IntoIterator<Item = StoredProviderUsageWindow>,
|
||
|
|
{
|
||
|
|
Self {
|
||
|
|
by_request_id: self.by_request_id,
|
||
detached_bodies: self.detached_bodies,
|
|||
provider_usage_windows: RwLock::new(items.into_iter().collect()),
|
|||
auth_api_keys: self.auth_api_keys,
|
|||
provider_catalog: self.provider_catalog,
|
|||
}
|
|||
|
|
}
|
||
|
|||
pub fn with_auth_api_key_repository(
|
|||
|
|
mut self,
|
||
|
|
repository: Arc<InMemoryAuthApiKeySnapshotRepository>,
|
||
|
|
) -> Self {
|
||
|
|
self.auth_api_keys = Some(repository);
|
||
|
|
self
|
||
|
|
}
|
||
|
|
|
||
pub fn with_provider_catalog_repository(
|
|||
|
|
mut self,
|
||
|
|
repository: Arc<InMemoryProviderCatalogReadRepository>,
|
||
|
|
) -> Self {
|
||
|
|
self.provider_catalog = Some(repository);
|
||
|
|
self
|
||
|
|
}
|
||
}
|
|||
|
|
|
||
fn usage_status_is_finalized(status: &str) -> bool {
|
|||
|
|
matches!(status, "completed" | "failed" | "cancelled")
|
||
|
|
}
|
||
|
|
|
||
|
|
fn usage_status_is_lifecycle(status: &str) -> bool {
|
||
|
|
matches!(status, "pending" | "streaming")
|
||
|
|
}
|
||
|
|
|
||
fn merge_usage_timing(existing: Option<u64>, incoming: Option<u64>) -> Option<u64> {
|
|||
|
|
match incoming {
|
||
|
|
Some(0) | None => existing.or(incoming),
|
||
|
|
Some(value) => Some(value),
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
fn accumulate_provider_api_key_usage_contribution(
|
|||
|
|
aggregates: &mut BTreeMap<String, ProviderApiKeyUsageContribution>,
|
||
|
|
contribution: ProviderApiKeyUsageContribution,
|
||
|
|
) {
|
||
|
|
let entry = aggregates
|
||
|
|
.entry(contribution.key_id.clone())
|
||
|
|
.or_insert_with(|| ProviderApiKeyUsageContribution {
|
||
|
|
key_id: contribution.key_id.clone(),
|
||
|
|
..ProviderApiKeyUsageContribution::default()
|
||
|
|
});
|
||
|
|
entry.request_count = entry
|
||
|
|
.request_count
|
||
|
|
.saturating_add(contribution.request_count);
|
||
|
|
entry.success_count = entry
|
||
|
|
.success_count
|
||
|
|
.saturating_add(contribution.success_count);
|
||
|
|
entry.error_count = entry.error_count.saturating_add(contribution.error_count);
|
||
|
|
entry.total_tokens = entry.total_tokens.saturating_add(contribution.total_tokens);
|
||
|
|
entry.total_cost_usd += contribution.total_cost_usd;
|
||
|
|
entry.total_response_time_ms = entry
|
||
|
|
.total_response_time_ms
|
||
|
|
.saturating_add(contribution.total_response_time_ms);
|
||
|
|
entry.last_used_at_unix_secs = match (
|
||
|
|
entry.last_used_at_unix_secs,
|
||
|
|
contribution.last_used_at_unix_secs,
|
||
|
|
) {
|
||
|
|
(Some(existing), Some(candidate)) => Some(existing.max(candidate)),
|
||
|
|
(None, Some(candidate)) => Some(candidate),
|
||
|
|
(existing, None) => existing,
|
||
|
|
};
|
||
|
|
}
|
||
|
|
|
||
fn accumulate_api_key_usage_contribution(
|
|||
|
|
aggregates: &mut BTreeMap<String, ApiKeyUsageContribution>,
|
||
|
|
contribution: ApiKeyUsageContribution,
|
||
|
|
) {
|
||
|
|
let entry = aggregates
|
||
|
|
.entry(contribution.api_key_id.clone())
|
||
|
|
.or_insert_with(|| ApiKeyUsageContribution {
|
||
|
|
api_key_id: contribution.api_key_id.clone(),
|
||
|
|
..ApiKeyUsageContribution::default()
|
||
|
|
});
|
||
|
|
entry.total_requests = entry
|
||
|
|
.total_requests
|
||
|
|
.saturating_add(contribution.total_requests);
|
||
|
|
entry.total_tokens = entry.total_tokens.saturating_add(contribution.total_tokens);
|
||
|
|
entry.total_cost_usd += contribution.total_cost_usd;
|
||
|
|
entry.last_used_at_unix_secs = match (
|
||
|
|
entry.last_used_at_unix_secs,
|
||
|
|
contribution.last_used_at_unix_secs,
|
||
|
|
) {
|
||
|
|
(Some(existing), Some(candidate)) => Some(existing.max(candidate)),
|
||
|
|
(None, Some(candidate)) => Some(candidate),
|
||
|
|
(existing, None) => existing,
|
||
|
|
};
|
||
|
|
}
|
||
|
|
|
||
fn usage_audit_client_family(item: &StoredRequestUsageAudit) -> Option<&str> {
|
|||
|
|
item.client_family
|
||
|
|
.as_deref()
|
||
|
|
.map(str::trim)
|
||
|
|
.filter(|value| !value.is_empty())
|
||
|
|
.or_else(|| usage_request_metadata_client_family(item.request_metadata.as_ref()))
|
||
|
|
}
|
||
|
|
|
||
|
|
fn usage_admin_unknown_label(value: &str) -> bool {
|
||
|
|
matches!(
|
||
|
|
value.trim().to_ascii_lowercase().as_str(),
|
||
|
|
"unknown" | "unknow"
|
||
|
|
)
|
||
|
|
}
|
||
|
|
|
||
|
|
fn usage_has_admin_unknown_model_or_provider(item: &StoredRequestUsageAudit) -> bool {
|
||
|
|
usage_admin_unknown_label(&item.model) || usage_admin_unknown_label(&item.provider_name)
|
||
|
|
}
|
||
|
|
|
||
fn usage_matches_list_query(item: &StoredRequestUsageAudit, query: &UsageAuditListQuery) -> bool {
|
|||
|
|
// The field is historically named `created_at_unix_ms`, but usage audit rows
|
||
|
|
// across gateway handlers, SQL repositories and tests are stored as epoch seconds.
|
||
|
|
if let Some(created_from_unix_secs) = query.created_from_unix_secs {
|
||
|
|
if item.created_at_unix_ms < created_from_unix_secs {
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if let Some(created_until_unix_secs) = query.created_until_unix_secs {
|
||
|
|
if item.created_at_unix_ms >= created_until_unix_secs {
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if let Some(user_id) = query.user_id.as_deref() {
|
||
|
|
if item.user_id.as_deref() != Some(user_id) {
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if let Some(provider_name) = query.provider_name.as_deref() {
|
||
|
|
if item.provider_name != provider_name {
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if let Some(model) = query.model.as_deref() {
|
||
|
|
if item.model != model {
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if let Some(api_format) = query.api_format.as_deref() {
|
||
|
|
if item.api_format.as_deref() != Some(api_format) {
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
}
|
||
if let Some(client_family) = query
|
|||
|
|
.client_family
|
||
|
|
.as_deref()
|
||
|
|
.map(str::trim)
|
||
|
|
.filter(|value| !value.is_empty())
|
||
|
|
{
|
||
|
|
if !usage_audit_client_family(item)
|
||
|
|
.is_some_and(|value| value.eq_ignore_ascii_case(client_family))
|
||
|
|
{
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if query.exclude_unknown_model_or_provider && usage_has_admin_unknown_model_or_provider(item) {
|
||
|
|
return false;
|
||
|
|
}
|
||
if let Some(statuses) = query.statuses.as_ref() {
|
|||
|
|
if !statuses.iter().any(|status| status == &item.status) {
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
}
|
||
if item
|
|||
|
|
.status_code
|
||
|
|
.is_some_and(|status_code| query.exclude_status_codes.contains(&status_code))
|
||
|
|
{
|
||
|
|
return false;
|
||
|
|
}
|
||
if let Some(is_stream) = query.is_stream {
|
|||
|
|
if item.is_stream != is_stream {
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if query.error_only
|
||
|
|
&& item.status != "failed"
|
||
|
|
&& item.status_code.unwrap_or_default() < 400
|
||
|
|
&& item
|
||
|
|
.error_message
|
||
|
|
.as_deref()
|
||
|
|
.map(str::trim)
|
||
|
|
.unwrap_or_default()
|
||
|
|
.is_empty()
|
||
|
|
{
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
|
||
|
|
true
|
||
|
|
}
|
||
|
|
|
||
fn usage_matches_keyword_search_query(
|
|||
|
|
item: &StoredRequestUsageAudit,
|
||
|
|
query: &UsageAuditKeywordSearchQuery,
|
||
|
|
) -> bool {
|
||
|
|
if let Some(created_from_unix_secs) = query.created_from_unix_secs {
|
||
|
|
if item.created_at_unix_ms < created_from_unix_secs {
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if let Some(created_until_unix_secs) = query.created_until_unix_secs {
|
||
|
|
if item.created_at_unix_ms >= created_until_unix_secs {
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if let Some(user_id) = query.user_id.as_deref() {
|
||
|
|
if item.user_id.as_deref() != Some(user_id) {
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if let Some(provider_name) = query.provider_name.as_deref() {
|
||
|
|
if item.provider_name != provider_name {
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if let Some(model) = query.model.as_deref() {
|
||
|
|
if item.model != model {
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if let Some(api_format) = query.api_format.as_deref() {
|
||
|
|
if item.api_format.as_deref() != Some(api_format) {
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
}
|
||
if let Some(client_family) = query
|
|||
|
|
.client_family
|
||
|
|
.as_deref()
|
||
|
|
.map(str::trim)
|
||
|
|
.filter(|value| !value.is_empty())
|
||
|
|
{
|
||
|
|
if !usage_audit_client_family(item)
|
||
|
|
.is_some_and(|value| value.eq_ignore_ascii_case(client_family))
|
||
|
|
{
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if query.exclude_unknown_model_or_provider && usage_has_admin_unknown_model_or_provider(item) {
|
||
|
|
return false;
|
||
|
|
}
|
||
if let Some(statuses) = query.statuses.as_ref() {
|
|||
|
|
if !statuses.iter().any(|status| status == &item.status) {
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
}
|
||
if item
|
|||
|
|
.status_code
|
||
|
|
.is_some_and(|status_code| query.exclude_status_codes.contains(&status_code))
|
||
|
|
{
|
||
|
|
return false;
|
||
|
|
}
|
||
if let Some(is_stream) = query.is_stream {
|
|||
|
|
if item.is_stream != is_stream {
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if query.error_only
|
||
|
|
&& item.status != "failed"
|
||
|
|
&& item.status_code.unwrap_or_default() < 400
|
||
|
|
&& item
|
||
|
|
.error_message
|
||
|
|
.as_deref()
|
||
|
|
.map(str::trim)
|
||
|
|
.unwrap_or_default()
|
||
|
|
.is_empty()
|
||
|
|
{
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
|
||
|
|
let model = item.model.to_ascii_lowercase();
|
||
|
|
let provider = item.provider_name.to_ascii_lowercase();
|
||
|
|
let legacy_username = item
|
||
|
|
.username
|
||
|
|
.as_deref()
|
||
|
|
.unwrap_or_default()
|
||
|
|
.to_ascii_lowercase();
|
||
|
|
let legacy_api_key_name = item
|
||
|
|
.api_key_name
|
||
|
|
.as_deref()
|
||
|
|
.unwrap_or_default()
|
||
|
|
.to_ascii_lowercase();
|
||
|
|
|
||
|
|
query.keywords.iter().enumerate().all(|(index, keyword)| {
|
||
|
|
let keyword = keyword.trim();
|
||
|
|
if keyword.is_empty() {
|
||
|
|
return true;
|
||
|
|
}
|
||
|
|
if model.contains(keyword) || provider.contains(keyword) {
|
||
|
|
return true;
|
||
|
|
}
|
||
|
|
if query.auth_user_reader_available {
|
||
|
|
let mut matched_user_ids = query
|
||
|
|
.matched_user_ids_by_keyword
|
||
|
|
.get(index)
|
||
|
|
.into_iter()
|
||
|
|
.flatten();
|
||
|
|
if item
|
||
|
|
.user_id
|
||
|
|
.as_ref()
|
||
|
|
.is_some_and(|user_id| matched_user_ids.any(|candidate| candidate == user_id))
|
||
|
|
{
|
||
|
|
return true;
|
||
|
|
}
|
||
|
|
} else if legacy_username.contains(keyword) {
|
||
|
|
return true;
|
||
|
|
}
|
||
|
|
if query.auth_api_key_reader_available {
|
||
|
|
let mut matched_ids = query
|
||
|
|
.matched_api_key_ids_by_keyword
|
||
|
|
.get(index)
|
||
|
|
.into_iter()
|
||
|
|
.flatten();
|
||
|
|
item.api_key_id
|
||
|
|
.as_ref()
|
||
|
|
.is_some_and(|api_key_id| matched_ids.any(|candidate| candidate == api_key_id))
|
||
|
|
} else {
|
||
|
|
legacy_api_key_name.contains(keyword)
|
||
|
|
}
|
||
|
|
}) && match query.username_keyword.as_deref().map(str::trim) {
|
||
|
|
Some(username_keyword) if !username_keyword.is_empty() => {
|
||
|
|
let username_keyword = username_keyword.to_ascii_lowercase();
|
||
|
|
if query.auth_user_reader_available {
|
||
|
|
item.user_id.as_ref().is_some_and(|user_id| {
|
||
|
|
query
|
||
|
|
.matched_user_ids_for_username
|
||
|
|
.iter()
|
||
|
|
.any(|candidate| candidate == user_id)
|
||
|
|
})
|
||
|
|
} else {
|
||
|
|
legacy_username.contains(&username_keyword)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
_ => true,
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
fn usage_matches_summary_query(
|
|||
|
|
item: &StoredRequestUsageAudit,
|
||
|
|
query: &UsageAuditSummaryQuery,
|
||
|
|
) -> bool {
|
||
|
|
if item.created_at_unix_ms < query.created_from_unix_secs
|
||
|
|
|| item.created_at_unix_ms >= query.created_until_unix_secs
|
||
|
|
{
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
if let Some(user_id) = query.user_id.as_deref() {
|
||
|
|
if item.user_id.as_deref() != Some(user_id) {
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if let Some(provider_name) = query.provider_name.as_deref() {
|
||
|
|
if item.provider_name != provider_name {
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if let Some(model) = query.model.as_deref() {
|
||
|
|
if item.model != model {
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
true
|
||
|
|
}
|
||
|
|
|
||
|
|
fn usage_matches_time_series_query(
|
||
|
|
item: &StoredRequestUsageAudit,
|
||
|
|
query: &UsageTimeSeriesQuery,
|
||
|
|
) -> bool {
|
||
|
|
if item.created_at_unix_ms < query.created_from_unix_secs
|
||
|
|
|| item.created_at_unix_ms >= query.created_until_unix_secs
|
||
|
|
{
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
if let Some(user_id) = query.user_id.as_deref() {
|
||
|
|
if item.user_id.as_deref() != Some(user_id) {
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if let Some(provider_name) = query.provider_name.as_deref() {
|
||
|
|
if item.provider_name != provider_name {
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if let Some(model) = query.model.as_deref() {
|
||
|
|
if item.model != model {
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
true
|
||
|
|
}
|
||
|
|
|
||
fn usage_matches_dashboard_summary_query(
|
|||
|
|
item: &StoredRequestUsageAudit,
|
||
|
|
query: &UsageDashboardSummaryQuery,
|
||
|
|
) -> bool {
|
||
|
|
if item.created_at_unix_ms < query.created_from_unix_secs
|
||
|
|
|| item.created_at_unix_ms >= query.created_until_unix_secs
|
||
|
|
|| matches!(item.status.as_str(), "pending" | "streaming")
|
||
|
|
|| matches!(item.provider_name.as_str(), "unknown" | "pending")
|
||
|
|
{
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
if let Some(user_id) = query.user_id.as_deref() {
|
||
|
|
if item.user_id.as_deref() != Some(user_id) {
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
true
|
||
|
|
}
|
||
|
|
|
||
|
|
fn usage_matches_dashboard_daily_breakdown_query(
|
||
|
|
item: &StoredRequestUsageAudit,
|
||
|
|
query: &UsageDashboardDailyBreakdownQuery,
|
||
|
|
) -> bool {
|
||
|
|
if item.created_at_unix_ms < query.created_from_unix_secs
|
||
|
|
|| item.created_at_unix_ms >= query.created_until_unix_secs
|
||
|
|
|| matches!(item.status.as_str(), "pending" | "streaming")
|
||
|
|
|| matches!(item.provider_name.as_str(), "unknown" | "pending")
|
||
|
|
{
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
if let Some(user_id) = query.user_id.as_deref() {
|
||
|
|
if item.user_id.as_deref() != Some(user_id) {
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
true
|
||
|
|
}
|
||
|
|
|
||
|
|
fn usage_matches_dashboard_provider_counts_query(
|
||
|
|
item: &StoredRequestUsageAudit,
|
||
|
|
query: &UsageDashboardProviderCountsQuery,
|
||
|
|
) -> bool {
|
||
|
|
if item.created_at_unix_ms < query.created_from_unix_secs
|
||
|
|
|| item.created_at_unix_ms >= query.created_until_unix_secs
|
||
|
|
|| matches!(item.status.as_str(), "pending" | "streaming")
|
||
|
|
|| matches!(item.provider_name.as_str(), "unknown" | "pending")
|
||
|
|
{
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
if let Some(user_id) = query.user_id.as_deref() {
|
||
|
|
if item.user_id.as_deref() != Some(user_id) {
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
true
|
||
|
|
}
|
||
|
|
|
||
|
|
fn usage_matches_breakdown_summary_query(
|
||
|
|
item: &StoredRequestUsageAudit,
|
||
|
|
query: &UsageBreakdownSummaryQuery,
|
||
|
|
) -> bool {
|
||
|
|
if item.created_at_unix_ms < query.created_from_unix_secs
|
||
|
|
|| item.created_at_unix_ms >= query.created_until_unix_secs
|
||
|
|
|| matches!(item.status.as_str(), "pending" | "streaming")
|
||
|
|
|| matches!(item.provider_name.as_str(), "unknown" | "pending")
|
||
|
|
{
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
if let Some(user_id) = query.user_id.as_deref() {
|
||
|
|
if item.user_id.as_deref() != Some(user_id) {
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
}
|
||
if let Some(provider_name) = query.provider_name.as_deref() {
|
|||
|
|
if item.provider_name != provider_name {
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
}
|
||
if let Some(model) = query.model.as_deref() {
|
|||
|
|
if item.model != model {
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if let Some(api_format) = query.api_format.as_deref() {
|
||
|
|
if item.api_format.as_deref() != Some(api_format) {
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
}
|
||
if item
|
|||
|
|
.status_code
|
||
|
|
.is_some_and(|status_code| query.exclude_status_codes.contains(&status_code))
|
||
|
|
{
|
||
|
|
return false;
|
||
|
|
}
|
||
match query.group_by {
|
|||
|
|
UsageBreakdownGroupBy::Model | UsageBreakdownGroupBy::Provider => true,
|
||
|
|
UsageBreakdownGroupBy::ApiFormat => item.api_format.is_some(),
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
fn usage_is_monitoring_error(item: &StoredRequestUsageAudit) -> bool {
|
||
|
|
let status = item.status.trim();
|
||
|
|
status.eq_ignore_ascii_case("error")
|
||
|
|
|| status.eq_ignore_ascii_case("failed")
|
||
|
|
|| item
|
||
|
|
.error_category
|
||
|
|
.as_deref()
|
||
|
|
.map(str::trim)
|
||
|
|
.is_some_and(|value| !value.is_empty())
|
||
|
|
|| (status.is_empty()
|
||
|
|
&& (item.status_code.unwrap_or_default() >= 400
|
||
|
|
|| item
|
||
|
|
.error_message
|
||
|
|
.as_deref()
|
||
|
|
.map(str::trim)
|
||
|
|
.is_some_and(|value| !value.is_empty())))
|
||
|
|
}
|
||
|
|
|
||
|
|
fn usage_matches_monitoring_error_count_query(
|
||
|
|
item: &StoredRequestUsageAudit,
|
||
|
|
query: &UsageMonitoringErrorCountQuery,
|
||
|
|
) -> bool {
|
||
|
|
item.created_at_unix_ms >= query.created_from_unix_secs
|
||
|
|
&& item.created_at_unix_ms < query.created_until_unix_secs
|
||
|
|
&& usage_is_monitoring_error(item)
|
||
|
|
}
|
||
|
|
|
||
|
|
fn usage_matches_monitoring_error_list_query(
|
||
|
|
item: &StoredRequestUsageAudit,
|
||
|
|
query: &UsageMonitoringErrorListQuery,
|
||
|
|
) -> bool {
|
||
|
|
item.created_at_unix_ms >= query.created_from_unix_secs
|
||
|
|
&& item.created_at_unix_ms < query.created_until_unix_secs
|
||
|
|
&& usage_is_monitoring_error(item)
|
||
|
|
}
|
||
|
|
|
||
|
|
fn usage_matches_error_distribution_query(
|
||
|
|
item: &StoredRequestUsageAudit,
|
||
|
|
query: &UsageErrorDistributionQuery,
|
||
|
|
) -> bool {
|
||
|
|
if item.created_at_unix_ms < query.created_from_unix_secs
|
||
|
|
|| item.created_at_unix_ms >= query.created_until_unix_secs
|
||
|
|
{
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
item.error_category
|
||
|
|
.as_deref()
|
||
|
|
.map(str::trim)
|
||
|
|
.is_some_and(|value| !value.is_empty())
|
||
|
|
}
|
||
|
|
|
||
|
|
fn usage_matches_performance_percentiles_query(
|
||
|
|
item: &StoredRequestUsageAudit,
|
||
|
|
query: &UsagePerformancePercentilesQuery,
|
||
|
|
) -> bool {
|
||
|
|
item.created_at_unix_ms >= query.created_from_unix_secs
|
||
|
|
&& item.created_at_unix_ms < query.created_until_unix_secs
|
||
|
|
&& item.status == "completed"
|
||
|
|
}
|
||
|
|
|
||
fn usage_reserved_provider_label(value: &str) -> bool {
|
|||
|
|
matches!(
|
||
|
|
value.trim().to_ascii_lowercase().as_str(),
|
||
|
|
"unknown" | "unknow" | "pending"
|
||
|
|
)
|
||
|
|
}
|
||
|
|
|
||
fn usage_provider_performance_identity(item: &StoredRequestUsageAudit) -> Option<(String, String)> {
|
|||
|
|
let provider_id = item.provider_id.as_deref()?.trim();
|
||
if provider_id.is_empty() || usage_reserved_provider_label(provider_id) {
|
|||
return None;
|
|||
|
|
}
|
||
|
|
let provider_name = item.provider_name.trim();
|
||
if usage_reserved_provider_label(provider_name) {
|
|||
return None;
|
|||
|
|
}
|
||
|
|
let display_name = if provider_name.is_empty() {
|
||
|
|
provider_id
|
||
|
|
} else {
|
||
|
|
provider_name
|
||
|
|
};
|
||
|
|
Some((provider_id.to_string(), display_name.to_string()))
|
||
|
|
}
|
||
|
|
|
||
|
|
fn usage_matches_provider_performance_query(
|
||
|
|
item: &StoredRequestUsageAudit,
|
||
|
|
query: &UsageProviderPerformanceQuery,
|
||
|
|
) -> Option<(String, String)> {
|
||
|
|
if item.created_at_unix_ms < query.created_from_unix_secs
|
||
|
|
|| item.created_at_unix_ms >= query.created_until_unix_secs
|
||
|
|
|| matches!(item.status.as_str(), "pending" | "streaming")
|
||
|
|
{
|
||
|
|
return None;
|
||
|
|
}
|
||
if let Some(provider_id) = query.provider_id.as_deref() {
|
|||
|
|
if item.provider_id.as_deref() != Some(provider_id) {
|
||
|
|
return None;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if let Some(model) = query.model.as_deref() {
|
||
|
|
if item.model != model {
|
||
|
|
return None;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if let Some(api_format) = query.api_format.as_deref() {
|
||
|
|
if item.api_format.as_deref() != Some(api_format) {
|
||
|
|
return None;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if let Some(endpoint_kind) = query.endpoint_kind.as_deref() {
|
||
|
|
if item.endpoint_kind.as_deref() != Some(endpoint_kind) {
|
||
|
|
return None;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if let Some(is_stream) = query.is_stream {
|
||
|
|
if item.is_stream != is_stream {
|
||
|
|
return None;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if let Some(has_format_conversion) = query.has_format_conversion {
|
||
|
|
if item.has_format_conversion != has_format_conversion {
|
||
|
|
return None;
|
||
|
|
}
|
||
|
|
}
|
||
usage_provider_performance_identity(item)
|
|||
|
|
}
|
||
|
|
|
||
fn usage_matches_cost_savings_query(
|
|||
|
|
item: &StoredRequestUsageAudit,
|
||
|
|
query: &UsageCostSavingsSummaryQuery,
|
||
|
|
) -> bool {
|
||
|
|
if item.created_at_unix_ms < query.created_from_unix_secs
|
||
|
|
|| item.created_at_unix_ms >= query.created_until_unix_secs
|
||
|
|
{
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
if let Some(user_id) = query.user_id.as_deref() {
|
||
|
|
if item.user_id.as_deref() != Some(user_id) {
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if let Some(provider_name) = query.provider_name.as_deref() {
|
||
|
|
if item.provider_name != provider_name {
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if let Some(model) = query.model.as_deref() {
|
||
|
|
if item.model != model {
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
true
|
||
|
|
}
|
||
|
|
|
||
|
|
fn usage_matches_cache_affinity_hit_summary_query(
|
||
|
|
item: &StoredRequestUsageAudit,
|
||
|
|
query: &UsageCacheAffinityHitSummaryQuery,
|
||
|
|
) -> bool {
|
||
|
|
if item.created_at_unix_ms < query.created_from_unix_secs
|
||
|
|
|| item.created_at_unix_ms >= query.created_until_unix_secs
|
||
|
|
|| item.status != "completed"
|
||
|
|
{
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
if let Some(user_id) = query.user_id.as_deref() {
|
||
|
|
if item.user_id.as_deref() != Some(user_id) {
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if let Some(api_key_id) = query.api_key_id.as_deref() {
|
||
|
|
if item.api_key_id.as_deref() != Some(api_key_id) {
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
true
|
||
|
|
}
|
||
|
|
|
||
|
|
fn usage_matches_cache_affinity_interval_query(
|
||
|
|
item: &StoredRequestUsageAudit,
|
||
|
|
query: &UsageCacheAffinityIntervalQuery,
|
||
|
|
) -> Option<String> {
|
||
|
|
if item.created_at_unix_ms < query.created_from_unix_secs
|
||
|
|
|| item.created_at_unix_ms >= query.created_until_unix_secs
|
||
|
|
|| item.status != "completed"
|
||
|
|
{
|
||
|
|
return None;
|
||
|
|
}
|
||
|
|
if let Some(user_id) = query.user_id.as_deref() {
|
||
|
|
if item.user_id.as_deref() != Some(user_id) {
|
||
|
|
return None;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if let Some(api_key_id) = query.api_key_id.as_deref() {
|
||
|
|
if item.api_key_id.as_deref() != Some(api_key_id) {
|
||
|
|
return None;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
match query.group_by {
|
||
|
|
UsageCacheAffinityIntervalGroupBy::User => item.user_id.clone(),
|
||
|
|
UsageCacheAffinityIntervalGroupBy::ApiKey => item.api_key_id.clone(),
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
fn usage_dashboard_local_date(
|
||
|
|
item: &StoredRequestUsageAudit,
|
||
|
|
tz_offset_minutes: i32,
|
||
|
|
) -> Option<String> {
|
||
|
|
let timestamp =
|
||
|
|
chrono::DateTime::<Utc>::from_timestamp(i64::try_from(item.created_at_unix_ms).ok()?, 0)?;
|
||
|
|
let local =
|
||
|
|
timestamp.checked_add_signed(chrono::Duration::minutes(i64::from(tz_offset_minutes)))?;
|
||
|
|
Some(local.date_naive().to_string())
|
||
|
|
}
|
||
|
|
|
||
|
|
fn usage_percentile_cont(values: &mut [u64], percentile: f64) -> Option<u64> {
|
||
|
|
if values.len() < 10 {
|
||
|
|
return None;
|
||
|
|
}
|
||
|
|
values.sort_unstable();
|
||
|
|
let position = percentile * (values.len().saturating_sub(1)) as f64;
|
||
|
|
let lower = position.floor() as usize;
|
||
|
|
let upper = position.ceil() as usize;
|
||
|
|
let lower_value = values[lower] as f64;
|
||
|
|
let upper_value = values[upper] as f64;
|
||
|
|
Some((lower_value + (upper_value - lower_value) * (position - lower as f64)).trunc() as u64)
|
||
|
|
}
|
||
|
|
|
||
fn usage_time_series_bucket_key(
|
|||
|
|
item: &StoredRequestUsageAudit,
|
||
|
|
granularity: UsageTimeSeriesGranularity,
|
||
|
|
tz_offset_minutes: i32,
|
||
|
|
) -> Option<String> {
|
||
|
|
let timestamp =
|
||
|
|
chrono::DateTime::<Utc>::from_timestamp(i64::try_from(item.created_at_unix_ms).ok()?, 0)?;
|
||
|
|
let local =
|
||
|
|
timestamp.checked_add_signed(chrono::Duration::minutes(i64::from(tz_offset_minutes)))?;
|
||
|
|
Some(match granularity {
|
||
|
|
UsageTimeSeriesGranularity::Day => local.date_naive().to_string(),
|
||
|
|
UsageTimeSeriesGranularity::Hour => local.format("%Y-%m-%dT%H:00:00+00:00").to_string(),
|
||
|
|
})
|
||
|
|
}
|
||
|
|
|
||
|
|
fn usage_matches_leaderboard_query(
|
||
|
|
item: &StoredRequestUsageAudit,
|
||
|
|
query: &UsageLeaderboardQuery,
|
||
|
|
) -> bool {
|
||
|
|
if item.created_at_unix_ms < query.created_from_unix_secs
|
||
|
|
|| item.created_at_unix_ms >= query.created_until_unix_secs
|
||
|
|
|| matches!(item.status.as_str(), "pending" | "streaming")
|
||
|
|
|| matches!(item.provider_name.as_str(), "unknown" | "pending")
|
||
|
|
{
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
if let Some(user_id) = query.user_id.as_deref() {
|
||
|
|
if item.user_id.as_deref() != Some(user_id) {
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if let Some(provider_name) = query.provider_name.as_deref() {
|
||
|
|
if item.provider_name != provider_name {
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if let Some(model) = query.model.as_deref() {
|
||
|
|
if item.model != model {
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
true
|
||
|
|
}
|
||
|
|
|
||
|
|
fn sort_usage_items(items: &mut [StoredRequestUsageAudit], newest_first: bool) {
|
||
|
|
items.sort_by(|left, right| {
|
||
|
|
let created_order = if newest_first {
|
||
|
|
right.created_at_unix_ms.cmp(&left.created_at_unix_ms)
|
||
|
|
} else {
|
||
|
|
left.created_at_unix_ms.cmp(&right.created_at_unix_ms)
|
||
|
|
};
|
||
|
|
if newest_first {
|
||
|
|
created_order.then_with(|| left.id.cmp(&right.id))
|
||
|
|
} else {
|
||
|
|
created_order.then_with(|| left.request_id.cmp(&right.request_id))
|
||
|
|
}
|
||
|
|
});
|
||
|
|
}
|
||
|
|
|
||
|
|
fn usage_cache_creation_tokens(item: &StoredRequestUsageAudit) -> u64 {
|
||
|
|
let classified = item
|
||
|
|
.cache_creation_ephemeral_5m_input_tokens
|
||
|
|
.saturating_add(item.cache_creation_ephemeral_1h_input_tokens);
|
||
|
|
if item.cache_creation_input_tokens == 0 && classified > 0 {
|
||
|
|
classified
|
||
|
|
} else {
|
||
|
|
item.cache_creation_input_tokens
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
#[derive(Clone, Copy, PartialEq, Eq)]
|
||
|
|
enum UsageApiFamily {
|
||
|
|
OpenAi,
|
||
|
|
Claude,
|
||
|
|
Gemini,
|
||
|
|
Unknown,
|
||
|
|
}
|
||
|
|
|
||
|
|
fn usage_api_family(api_format: Option<&str>) -> UsageApiFamily {
|
||
|
|
let Some(api_format) = api_format else {
|
||
|
|
return UsageApiFamily::Unknown;
|
||
|
|
};
|
||
|
|
let family = api_format
|
||
|
|
.split(':')
|
||
|
|
.next()
|
||
|
|
.unwrap_or_default()
|
||
|
|
.trim()
|
||
|
|
.to_ascii_lowercase();
|
||
|
|
match family.as_str() {
|
||
|
|
"openai" => UsageApiFamily::OpenAi,
|
||
|
|
"claude" | "anthropic" => UsageApiFamily::Claude,
|
||
|
|
"gemini" | "google" => UsageApiFamily::Gemini,
|
||
|
|
_ => UsageApiFamily::Unknown,
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
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_creation_tokens <= 0 && cache_read_tokens <= 0 {
|
|||
return input_tokens;
|
|||
|
|
}
|
||
|
|
|
||
|
|
match usage_api_family(api_format) {
|
||
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,
|
|||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
fn normalize_usage_total_input_context(
|
||
|
|
api_format: Option<&str>,
|
||
|
|
input_tokens: i64,
|
||
|
|
cache_creation_tokens: i64,
|
||
|
|
cache_read_tokens: i64,
|
||
|
|
) -> i64 {
|
||
|
|
let normalized_input_tokens = input_tokens.max(0);
|
||
|
|
let normalized_cache_creation_tokens = cache_creation_tokens.max(0);
|
||
|
|
let normalized_cache_read_tokens = cache_read_tokens.max(0);
|
||
|
|
|
||
|
|
let fresh_input_tokens = match usage_api_family(api_format) {
|
||
|
|
UsageApiFamily::Claude => {
|
||
|
|
normalized_input_tokens.saturating_add(normalized_cache_creation_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 => {
|
||
|
|
if normalized_cache_creation_tokens > 0 {
|
||
|
|
normalized_input_tokens.saturating_add(normalized_cache_creation_tokens)
|
||
|
|
} else {
|
||
|
|
normalized_input_tokens
|
||
|
|
}
|
||
|
|
}
|
||
|
|
};
|
||
|
|
|
||
|
|
fresh_input_tokens.saturating_add(normalized_cache_read_tokens)
|
||
|
|
}
|
||
|
|
|
||
|
|
fn usage_total_input_context(item: &StoredRequestUsageAudit) -> u64 {
|
||
|
|
let api_format = item
|
||
|
|
.endpoint_api_format
|
||
|
|
.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_total_input_context(
|
||
|
|
api_format,
|
||
|
|
input_tokens,
|
||
|
|
cache_creation_tokens,
|
||
|
|
cache_read_tokens,
|
||
|
|
) as u64
|
||
|
|
}
|
||
|
|
|
||
|
|
fn usage_effective_input_tokens(item: &StoredRequestUsageAudit) -> u64 {
|
||
|
|
let api_format = item
|
||
|
|
.endpoint_api_format
|
||
|
|
.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_creation_tokens,
|
||
|
|
cache_read_tokens,
|
||
|
|
) as u64
|
||
}
|
|||
|
|
|
||
fn usage_total_tokens(item: &StoredRequestUsageAudit) -> u64 {
|
|||
|
|
usage_effective_input_tokens(item)
|
||
|
|
.saturating_add(item.output_tokens)
|
||
|
|
.saturating_add(usage_cache_creation_tokens(item))
|
||
|
|
.saturating_add(item.cache_read_input_tokens)
|
||
|
|
}
|
||
|
|
|
||
fn usage_is_success(item: &StoredRequestUsageAudit) -> bool {
|
|||
|
|
matches!(
|
||
|
|
item.status.as_str(),
|
||
|
|
"completed" | "success" | "ok" | "billed" | "settled"
|
||
|
|
) && item.status_code.is_none_or(|code| code < 400)
|
||
|
|
}
|
||
|
|
|
||
fn usage_output_tps_uses_generation_time(item: &StoredRequestUsageAudit) -> bool {
|
|||
|
|
item.request_metadata
|
||
|
|
.as_ref()
|
||
|
|
.and_then(Value::as_object)
|
||
.and_then(|metadata| metadata.get(UPSTREAM_IS_STREAM_KEY))
|
|||
.and_then(Value::as_bool)
|
|||
|
|
.unwrap_or(item.is_stream)
|
||
|
|
}
|
||
|
|
|
||
|
|
fn usage_output_tps_duration_ms(item: &StoredRequestUsageAudit) -> Option<u64> {
|
||
|
|
let response_time_ms = item.response_time_ms?;
|
||
|
|
if response_time_ms == 0 {
|
||
|
|
return None;
|
||
|
|
}
|
||
|
|
|
||
|
|
if !usage_output_tps_uses_generation_time(item) {
|
||
|
|
return Some(response_time_ms);
|
||
|
|
}
|
||
|
|
|
||
|
|
let first_byte_time_ms = item.first_byte_time_ms?;
|
||
|
|
if first_byte_time_ms >= response_time_ms {
|
||
|
|
return None;
|
||
|
|
}
|
||
|
|
Some(response_time_ms - first_byte_time_ms)
|
||
|
|
}
|
||
|
|
|
||
fn usage_provider_display_name(item: &StoredRequestUsageAudit) -> Option<String> {
|
|||
|
|
let provider_name = item.provider_name.trim();
|
||
if provider_name.is_empty() || usage_reserved_provider_label(provider_name) {
|
|||
None
|
|||
|
|
} else {
|
||
Some(provider_name.to_string())
|
|||
}
|
|||
|
|
}
|
||
|
|
|
||
fn usage_provider_id(item: &StoredRequestUsageAudit) -> Option<String> {
|
|||
|
|
let provider_id = item.provider_id.as_deref()?.trim();
|
||
|
|
if provider_id.is_empty() || usage_reserved_provider_label(provider_id) {
|
||
|
|
None
|
||
|
|
} else {
|
||
|
|
Some(provider_id.to_string())
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
fn usage_provider_aggregation_identity(
|
||
|
|
item: &StoredRequestUsageAudit,
|
||
|
|
) -> Option<(String, Option<String>, String)> {
|
||
|
|
let display_name = usage_provider_display_name(item);
|
||
|
|
if let Some(provider_id) = usage_provider_id(item) {
|
||
|
|
return Some((provider_id, display_name, "provider_id".to_string()));
|
||
|
|
}
|
||
|
|
let display_name = display_name?;
|
||
|
|
Some((
|
||
|
|
display_name.clone(),
|
||
|
|
Some(display_name),
|
||
|
|
"legacy_name".to_string(),
|
||
|
|
))
|
||
|
|
}
|
||
|
|
|
||
#[async_trait]
|
|||
|
|
impl UsageReadRepository for InMemoryUsageReadRepository {
|
||
async fn find_by_id(
|
|||
|
|
&self,
|
||
|
|
id: &str,
|
||
|
|
) -> Result<Option<StoredRequestUsageAudit>, DataLayerError> {
|
||
|
|
Ok(self
|
||
|
|
.by_request_id
|
||
|
|
.read()
|
||
|
|
.expect("usage repository lock")
|
||
|
|
.values()
|
||
|
|
.find(|item| item.id == id)
|
||
|
|
.cloned())
|
||
|
|
}
|
||
|
|
|
||
async fn list_by_ids(
|
|||
|
|
&self,
|
||
|
|
ids: &[String],
|
||
|
|
) -> Result<Vec<StoredRequestUsageAudit>, DataLayerError> {
|
||
|
|
if ids.is_empty() {
|
||
|
|
return Ok(Vec::new());
|
||
|
|
}
|
||
|
|
let requested_ids = ids
|
||
|
|
.iter()
|
||
|
|
.cloned()
|
||
|
|
.collect::<std::collections::BTreeSet<_>>();
|
||
|
|
Ok(self
|
||
|
|
.by_request_id
|
||
|
|
.read()
|
||
|
|
.expect("usage repository lock")
|
||
|
|
.values()
|
||
|
|
.filter(|item| requested_ids.contains(&item.id))
|
||
|
|
.cloned()
|
||
|
|
.collect())
|
||
|
|
}
|
||
|
|
|
||
async fn find_by_request_id(
|
|||
|
|
&self,
|
||
|
|
request_id: &str,
|
||
|
|
) -> Result<Option<StoredRequestUsageAudit>, DataLayerError> {
|
||
|
|
Ok(self
|
||
|
|
.by_request_id
|
||
|
|
.read()
|
||
|
|
.expect("usage repository lock")
|
||
|
|
.get(request_id)
|
||
|
|
.cloned())
|
||
|
|
}
|
||
|
|||
async fn resolve_body_ref(&self, body_ref: &str) -> Result<Option<Value>, DataLayerError> {
|
|||
|
|
if let Some(value) = self
|
||
|
|
.detached_bodies
|
||
|
|
.read()
|
||
|
|
.expect("usage repository lock")
|
||
|
|
.get(body_ref)
|
||
|
|
.cloned()
|
||
|
|
{
|
||
|
|
return Ok(Some(value));
|
||
|
|
}
|
||
|
|
let Some((request_id, field)) = parse_usage_body_ref(body_ref) else {
|
||
|
|
return Ok(None);
|
||
|
|
};
|
||
|
|
let usage = self
|
||
|
|
.by_request_id
|
||
|
|
.read()
|
||
|
|
.expect("usage repository lock")
|
||
|
|
.get(&request_id)
|
||
|
|
.cloned();
|
||
|
|
Ok(usage.and_then(|usage| match field {
|
||
|
|
UsageBodyField::RequestBody => usage.request_body,
|
||
|
|
UsageBodyField::ProviderRequestBody => usage.provider_request_body,
|
||
|
|
UsageBodyField::ResponseBody => usage.response_body,
|
||
|
|
UsageBodyField::ClientResponseBody => usage.client_response_body,
|
||
|
|
}))
|
||
|
|
}
|
||
|
|
|
||
async fn list_usage_audits(
|
|||
|
|
&self,
|
||
|
|
query: &UsageAuditListQuery,
|
||
|
|
) -> Result<Vec<StoredRequestUsageAudit>, DataLayerError> {
|
||
|
|
let mut items: Vec<_> = self
|
||
|
|
.by_request_id
|
||
|
|
.read()
|
||
|
|
.expect("usage repository lock")
|
||
|
|
.values()
|
||
.filter(|item| usage_matches_list_query(item, query))
|
|||
.cloned()
|
|||
|
|
.collect();
|
||
sort_usage_items(&mut items, query.newest_first);
|
|||
|
|
if let Some(offset) = query.offset {
|
||
|
|
if offset >= items.len() {
|
||
|
|
items.clear();
|
||
|
|
} else {
|
||
|
|
items.drain(..offset);
|
||
|
|
}
|
||
|
|
}
|
||
if let Some(limit) = query.limit {
|
|||
|
|
items.truncate(limit);
|
||
|
|
}
|
||
Ok(items)
|
|||
|
|
}
|
||
|
|
|
||
async fn list_usage_audits_by_keyword_search(
|
|||
|
|
&self,
|
||
|
|
query: &UsageAuditKeywordSearchQuery,
|
||
|
|
) -> Result<Vec<StoredRequestUsageAudit>, DataLayerError> {
|
||
|
|
let mut items: Vec<_> = self
|
||
|
|
.by_request_id
|
||
|
|
.read()
|
||
|
|
.expect("usage repository lock")
|
||
|
|
.values()
|
||
|
|
.filter(|item| usage_matches_keyword_search_query(item, query))
|
||
|
|
.cloned()
|
||
|
|
.collect();
|
||
|
|
sort_usage_items(&mut items, query.newest_first);
|
||
|
|
if let Some(offset) = query.offset {
|
||
|
|
if offset >= items.len() {
|
||
|
|
items.clear();
|
||
|
|
} else {
|
||
|
|
items.drain(..offset);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if let Some(limit) = query.limit {
|
||
|
|
items.truncate(limit);
|
||
|
|
}
|
||
|
|
Ok(items)
|
||
|
|
}
|
||
|
|
|
||
async fn count_usage_audits(&self, query: &UsageAuditListQuery) -> Result<u64, DataLayerError> {
|
|||
|
|
Ok(self
|
||
|
|
.by_request_id
|
||
|
|
.read()
|
||
|
|
.expect("usage repository lock")
|
||
|
|
.values()
|
||
|
|
.filter(|item| usage_matches_list_query(item, query))
|
||
|
|
.count() as u64)
|
||
|
|
}
|
||
|
|
|
||
async fn count_usage_audits_by_keyword_search(
|
|||
|
|
&self,
|
||
|
|
query: &UsageAuditKeywordSearchQuery,
|
||
|
|
) -> Result<u64, DataLayerError> {
|
||
|
|
Ok(self
|
||
|
|
.by_request_id
|
||
|
|
.read()
|
||
|
|
.expect("usage repository lock")
|
||
|
|
.values()
|
||
|
|
.filter(|item| usage_matches_keyword_search_query(item, query))
|
||
|
|
.count() as u64)
|
||
|
|
}
|
||
|
|
|
||
async fn aggregate_usage_audits(
|
|||
|
|
&self,
|
||
|
|
query: &UsageAuditAggregationQuery,
|
||
|
|
) -> Result<Vec<StoredUsageAuditAggregation>, DataLayerError> {
|
||
|
|
#[derive(Default)]
|
||
|
|
struct AggregateBucket {
|
||
|
|
display_name: Option<String>,
|
||
|
|
secondary_name: Option<String>,
|
||
|
|
request_count: u64,
|
||
|
|
total_tokens: u64,
|
||
|
|
output_tokens: u64,
|
||
|
|
effective_input_tokens: u64,
|
||
|
|
total_input_context: u64,
|
||
|
|
cache_creation_tokens: u64,
|
||
|
|
cache_creation_ephemeral_5m_tokens: u64,
|
||
|
|
cache_creation_ephemeral_1h_tokens: u64,
|
||
|
|
cache_read_tokens: u64,
|
||
|
|
total_cost_usd: f64,
|
||
|
|
actual_total_cost_usd: f64,
|
||
|
|
response_time_ms_sum: u64,
|
||
|
|
success_count: u64,
|
||
|
|
}
|
||
|
|
|
||
|
|
let mut grouped: BTreeMap<String, AggregateBucket> = BTreeMap::new();
|
||
|
|
for item in self
|
||
|
|
.by_request_id
|
||
|
|
.read()
|
||
|
|
.expect("usage repository lock")
|
||
|
|
.values()
|
||
|
|
{
|
||
|
|
if item.created_at_unix_ms < query.created_from_unix_secs
|
||
|
|
|| item.created_at_unix_ms >= query.created_until_unix_secs
|
||
|
|
|| matches!(item.status.as_str(), "pending" | "streaming")
|
||
|| (query.exclude_reserved_provider_labels
|
|||
&& usage_provider_aggregation_identity(item).is_none())
|
|||
{
|
|||
|
|
continue;
|
||
|
|
}
|
||
|
|
|
||
let provider_identity =
|
|||
if matches!(query.group_by, UsageAuditAggregationGroupBy::Provider) {
|
|||
match usage_provider_aggregation_identity(item) {
|
|||
Some(value) => Some(value),
|
|||
|
|
None => continue,
|
||
|
|
}
|
||
|
|
} else {
|
||
|
|
None
|
||
|
|
};
|
||
|
|
|
||
let group_key = match query.group_by {
|
|||
|
|
UsageAuditAggregationGroupBy::Model => item.model.clone(),
|
||
UsageAuditAggregationGroupBy::Provider => provider_identity
|
|||
|
|
.as_ref()
|
||
|
|
.expect("provider identity is set for provider aggregation")
|
||
|
|
.0
|
||
|
|
.clone(),
|
||
UsageAuditAggregationGroupBy::ApiFormat => item
|
|||
|
|
.api_format
|
||
|
|
.clone()
|
||
|
|
.unwrap_or_else(|| "unknown".to_string()),
|
||
|
|
UsageAuditAggregationGroupBy::User => match item.user_id.clone() {
|
||
|
|
Some(value) => value,
|
||
|
|
None => continue,
|
||
|
|
},
|
||
|
|
};
|
||
|
|
let bucket = grouped.entry(group_key).or_default();
|
||
|
|
if matches!(query.group_by, UsageAuditAggregationGroupBy::Provider)
|
||
|
|
&& (bucket.display_name.is_none()
|
||
|
|
|| bucket.display_name.as_deref() == Some("Unknown"))
|
||
|
|
{
|
||
bucket.display_name = provider_identity
|
|||
|
|
.as_ref()
|
||
|
|
.and_then(|(_, display_name, _)| display_name.clone());
|
||
|
|
}
|
||
|
|
if matches!(query.group_by, UsageAuditAggregationGroupBy::Provider)
|
||
|
|
&& (bucket.secondary_name.is_none()
|
||
|
|
|| bucket.secondary_name.as_deref() == Some("legacy_name"))
|
||
|
|
{
|
||
|
|
bucket.secondary_name = provider_identity
|
||
|
|
.as_ref()
|
||
|
|
.map(|(_, _, identity_source)| identity_source.clone());
|
||
}
|
|||
|
|
bucket.request_count = bucket.request_count.saturating_add(1);
|
||
|
|
bucket.total_tokens = bucket.total_tokens.saturating_add(item.total_tokens);
|
||
|
|
bucket.output_tokens = bucket.output_tokens.saturating_add(item.output_tokens);
|
||
|
|
bucket.effective_input_tokens = bucket
|
||
|
|
.effective_input_tokens
|
||
|
|
.saturating_add(usage_effective_input_tokens(item));
|
||
|
|
bucket.total_input_context = bucket
|
||
|
|
.total_input_context
|
||
|
|
.saturating_add(usage_total_input_context(item));
|
||
|
|
bucket.cache_creation_tokens = bucket
|
||
|
|
.cache_creation_tokens
|
||
|
|
.saturating_add(usage_cache_creation_tokens(item));
|
||
|
|
bucket.cache_creation_ephemeral_5m_tokens = bucket
|
||
|
|
.cache_creation_ephemeral_5m_tokens
|
||
|
|
.saturating_add(item.cache_creation_ephemeral_5m_input_tokens);
|
||
|
|
bucket.cache_creation_ephemeral_1h_tokens = bucket
|
||
|
|
.cache_creation_ephemeral_1h_tokens
|
||
|
|
.saturating_add(item.cache_creation_ephemeral_1h_input_tokens);
|
||
|
|
bucket.cache_read_tokens = bucket
|
||
|
|
.cache_read_tokens
|
||
|
|
.saturating_add(item.cache_read_input_tokens);
|
||
|
|
bucket.total_cost_usd += item.total_cost_usd;
|
||
|
|
bucket.actual_total_cost_usd += item.actual_total_cost_usd;
|
||
|
|
bucket.response_time_ms_sum = bucket
|
||
|
|
.response_time_ms_sum
|
||
|
|
.saturating_add(item.response_time_ms.unwrap_or_default());
|
||
|
|
bucket.success_count = bucket
|
||
|
|
.success_count
|
||
|
|
.saturating_add(if usage_is_success(item) { 1 } else { 0 });
|
||
|
|
}
|
||
|
|
|
||
|
|
let mut items = grouped
|
||
|
|
.into_iter()
|
||
|
|
.map(|(group_key, bucket)| StoredUsageAuditAggregation {
|
||
|
|
group_key,
|
||
|
|
display_name: bucket.display_name,
|
||
|
|
secondary_name: bucket.secondary_name,
|
||
|
|
request_count: bucket.request_count,
|
||
|
|
total_tokens: bucket.total_tokens,
|
||
|
|
output_tokens: bucket.output_tokens,
|
||
|
|
effective_input_tokens: bucket.effective_input_tokens,
|
||
|
|
total_input_context: bucket.total_input_context,
|
||
|
|
cache_creation_tokens: bucket.cache_creation_tokens,
|
||
|
|
cache_creation_ephemeral_5m_tokens: bucket.cache_creation_ephemeral_5m_tokens,
|
||
|
|
cache_creation_ephemeral_1h_tokens: bucket.cache_creation_ephemeral_1h_tokens,
|
||
|
|
cache_read_tokens: bucket.cache_read_tokens,
|
||
|
|
total_cost_usd: bucket.total_cost_usd,
|
||
|
|
actual_total_cost_usd: bucket.actual_total_cost_usd,
|
||
|
|
avg_response_time_ms: match query.group_by {
|
||
|
|
UsageAuditAggregationGroupBy::Provider
|
||
|
|
| UsageAuditAggregationGroupBy::ApiFormat => {
|
||
|
|
Some(if bucket.request_count == 0 {
|
||
|
|
0.0
|
||
|
|
} else {
|
||
|
|
bucket.response_time_ms_sum as f64 / bucket.request_count as f64
|
||
|
|
})
|
||
|
|
}
|
||
|
|
_ => None,
|
||
|
|
},
|
||
|
|
success_count: match query.group_by {
|
||
|
|
UsageAuditAggregationGroupBy::Provider => Some(bucket.success_count),
|
||
|
|
_ => None,
|
||
|
|
},
|
||
|
|
})
|
||
|
|
.collect::<Vec<_>>();
|
||
|
|
|
||
|
|
items.sort_by(|left, right| {
|
||
|
|
right
|
||
|
|
.request_count
|
||
|
|
.cmp(&left.request_count)
|
||
|
|
.then_with(|| left.group_key.cmp(&right.group_key))
|
||
|
|
});
|
||
|
|
items.truncate(query.limit);
|
||
|
|
Ok(items)
|
||
|
|
}
|
||
|
|
|
||
|
|
async fn summarize_usage_audits(
|
||
|
|
&self,
|
||
|
|
query: &UsageAuditSummaryQuery,
|
||
|
|
) -> Result<StoredUsageAuditSummary, DataLayerError> {
|
||
|
|
let mut summary = StoredUsageAuditSummary::default();
|
||
|
|
for item in self
|
||
|
|
.by_request_id
|
||
|
|
.read()
|
||
|
|
.expect("usage repository lock")
|
||
|
|
.values()
|
||
|
|
{
|
||
|
|
if !usage_matches_summary_query(item, query) {
|
||
|
|
continue;
|
||
|
|
}
|
||
|
|
summary.total_requests = summary.total_requests.saturating_add(1);
|
||
|
|
summary.input_tokens = summary.input_tokens.saturating_add(item.input_tokens);
|
||
|
|
summary.output_tokens = summary.output_tokens.saturating_add(item.output_tokens);
|
||
|
|
summary.recorded_total_tokens = summary
|
||
|
|
.recorded_total_tokens
|
||
|
|
.saturating_add(item.total_tokens);
|
||
|
|
summary.cache_creation_tokens = summary
|
||
|
|
.cache_creation_tokens
|
||
|
|
.saturating_add(usage_cache_creation_tokens(item));
|
||
|
|
summary.cache_creation_ephemeral_5m_tokens = summary
|
||
|
|
.cache_creation_ephemeral_5m_tokens
|
||
|
|
.saturating_add(item.cache_creation_ephemeral_5m_input_tokens);
|
||
|
|
summary.cache_creation_ephemeral_1h_tokens = summary
|
||
|
|
.cache_creation_ephemeral_1h_tokens
|
||
|
|
.saturating_add(item.cache_creation_ephemeral_1h_input_tokens);
|
||
|
|
summary.cache_read_tokens = summary
|
||
|
|
.cache_read_tokens
|
||
|
|
.saturating_add(item.cache_read_input_tokens);
|
||
|
|
summary.total_cost_usd += item.total_cost_usd;
|
||
|
|
summary.actual_total_cost_usd += item.actual_total_cost_usd;
|
||
|
|
summary.cache_creation_cost_usd += item.cache_creation_cost_usd;
|
||
|
|
summary.cache_read_cost_usd += item.cache_read_cost_usd;
|
||
|
|
summary.total_response_time_ms += item.response_time_ms.unwrap_or(0) as f64;
|
||
|
|
if item.status_code.is_some_and(|value| value >= 400) || item.error_message.is_some() {
|
||
|
|
summary.error_requests = summary.error_requests.saturating_add(1);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
Ok(summary)
|
||
|
|
}
|
||
|
|
|
||
async fn summarize_usage_cache_hit_summary(
|
|||
|
|
&self,
|
||
|
|
query: &UsageCacheHitSummaryQuery,
|
||
|
|
) -> Result<StoredUsageCacheHitSummary, DataLayerError> {
|
||
|
|
let mut summary = StoredUsageCacheHitSummary::default();
|
||
|
|
for item in self
|
||
|
|
.by_request_id
|
||
|
|
.read()
|
||
|
|
.expect("usage repository lock")
|
||
|
|
.values()
|
||
|
|
{
|
||
|
|
if item.created_at_unix_ms < query.created_from_unix_secs
|
||
|
|
|| item.created_at_unix_ms >= query.created_until_unix_secs
|
||
|
|
{
|
||
|
|
continue;
|
||
|
|
}
|
||
|
|
if let Some(user_id) = query.user_id.as_deref() {
|
||
|
|
if item.user_id.as_deref() != Some(user_id) {
|
||
|
|
continue;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
summary.total_requests = summary.total_requests.saturating_add(1);
|
||
|
|
if item.cache_read_input_tokens > 0 {
|
||
|
|
summary.cache_hit_requests = summary.cache_hit_requests.saturating_add(1);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
Ok(summary)
|
||
|
|
}
|
||
|
|
|
||
|
|
async fn summarize_usage_settled_cost(
|
||
|
|
&self,
|
||
|
|
query: &UsageSettledCostSummaryQuery,
|
||
|
|
) -> Result<StoredUsageSettledCostSummary, DataLayerError> {
|
||
|
|
let mut summary = StoredUsageSettledCostSummary::default();
|
||
|
|
for item in self
|
||
|
|
.by_request_id
|
||
|
|
.read()
|
||
|
|
.expect("usage repository lock")
|
||
|
|
.values()
|
||
|
|
{
|
||
|
|
if item.created_at_unix_ms < query.created_from_unix_secs
|
||
|
|
|| item.created_at_unix_ms >= query.created_until_unix_secs
|
||
|
|
{
|
||
|
|
continue;
|
||
|
|
}
|
||
|
|
if let Some(user_id) = query.user_id.as_deref() {
|
||
|
|
if item.user_id.as_deref() != Some(user_id) {
|
||
|
|
continue;
|
||
|
|
}
|
||
|
|
}
|
||
if let Some(api_key_id) = query.api_key_id.as_deref() {
|
|||
|
|
if item.api_key_id.as_deref() != Some(api_key_id) {
|
||
|
|
continue;
|
||
|
|
}
|
||
|
|
}
|
||
if item.billing_status != "settled" || item.total_cost_usd <= 0.0 {
|
|||
|
|
continue;
|
||
|
|
}
|
||
|
|
|
||
|
|
summary.total_cost_usd += item.total_cost_usd;
|
||
|
|
summary.total_requests = summary.total_requests.saturating_add(1);
|
||
|
|
summary.input_tokens = summary.input_tokens.saturating_add(item.input_tokens);
|
||
|
|
summary.output_tokens = summary.output_tokens.saturating_add(item.output_tokens);
|
||
|
|
summary.cache_creation_tokens = summary
|
||
|
|
.cache_creation_tokens
|
||
|
|
.saturating_add(item.cache_creation_input_tokens);
|
||
|
|
summary.cache_read_tokens = summary
|
||
|
|
.cache_read_tokens
|
||
|
|
.saturating_add(item.cache_read_input_tokens);
|
||
|
|
if let Some(finalized_at_unix_secs) = item.finalized_at_unix_secs {
|
||
|
|
summary.first_finalized_at_unix_secs = Some(
|
||
|
|
summary
|
||
|
|
.first_finalized_at_unix_secs
|
||
|
|
.map(|value| value.min(finalized_at_unix_secs))
|
||
|
|
.unwrap_or(finalized_at_unix_secs),
|
||
|
|
);
|
||
|
|
summary.last_finalized_at_unix_secs = Some(
|
||
|
|
summary
|
||
|
|
.last_finalized_at_unix_secs
|
||
|
|
.map(|value| value.max(finalized_at_unix_secs))
|
||
|
|
.unwrap_or(finalized_at_unix_secs),
|
||
|
|
);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
Ok(summary)
|
||
|
|
}
|
||
|
|
|
||
|
|
async fn summarize_dashboard_usage(
|
||
|
|
&self,
|
||
|
|
query: &UsageDashboardSummaryQuery,
|
||
|
|
) -> Result<StoredUsageDashboardSummary, DataLayerError> {
|
||
|
|
let mut summary = StoredUsageDashboardSummary::default();
|
||
|
|
for item in self
|
||
|
|
.by_request_id
|
||
|
|
.read()
|
||
|
|
.expect("usage repository lock")
|
||
|
|
.values()
|
||
|
|
{
|
||
|
|
if !usage_matches_dashboard_summary_query(item, query) {
|
||
|
|
continue;
|
||
|
|
}
|
||
|
|
summary.total_requests = summary.total_requests.saturating_add(1);
|
||
|
|
summary.input_tokens = summary.input_tokens.saturating_add(item.input_tokens);
|
||
|
|
summary.effective_input_tokens = summary
|
||
|
|
.effective_input_tokens
|
||
|
|
.saturating_add(usage_effective_input_tokens(item));
|
||
|
|
summary.output_tokens = summary.output_tokens.saturating_add(item.output_tokens);
|
||
summary.total_tokens = summary
|
|||
|
|
.total_tokens
|
||
|
|
.saturating_add(usage_total_tokens(item));
|
||
summary.cache_creation_tokens = summary
|
|||
|
|
.cache_creation_tokens
|
||
|
|
.saturating_add(usage_cache_creation_tokens(item));
|
||
|
|
summary.cache_read_tokens = summary
|
||
|
|
.cache_read_tokens
|
||
|
|
.saturating_add(item.cache_read_input_tokens);
|
||
|
|
summary.total_input_context = summary
|
||
|
|
.total_input_context
|
||
|
|
.saturating_add(usage_total_input_context(item));
|
||
|
|
summary.cache_creation_cost_usd += item.cache_creation_cost_usd;
|
||
|
|
summary.cache_read_cost_usd += item.cache_read_cost_usd;
|
||
|
|
summary.total_cost_usd += item.total_cost_usd;
|
||
|
|
summary.actual_total_cost_usd += item.actual_total_cost_usd;
|
||
|
|
if item.status_code.is_some_and(|value| value >= 400)
|
||
|
|
|| item.status.eq_ignore_ascii_case("failed")
|
||
|
|
{
|
||
|
|
summary.error_requests = summary.error_requests.saturating_add(1);
|
||
|
|
}
|
||
|
|
if let Some(response_time_ms) = item.response_time_ms {
|
||
|
|
summary.response_time_sum_ms += response_time_ms as f64;
|
||
|
|
summary.response_time_samples = summary.response_time_samples.saturating_add(1);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
Ok(summary)
|
||
|
|
}
|
||
|
|
|
||
|
|
async fn list_dashboard_daily_breakdown(
|
||
|
|
&self,
|
||
|
|
query: &UsageDashboardDailyBreakdownQuery,
|
||
|
|
) -> Result<Vec<StoredUsageDashboardDailyBreakdownRow>, DataLayerError> {
|
||
|
|
#[derive(Default)]
|
||
|
|
struct DashboardDailyBucket {
|
||
|
|
requests: u64,
|
||
|
|
total_tokens: u64,
|
||
|
|
total_cost_usd: f64,
|
||
|
|
response_time_sum_ms: f64,
|
||
|
|
response_time_samples: u64,
|
||
|
|
}
|
||
|
|
|
||
|
|
let mut grouped = BTreeMap::<(String, String, String), DashboardDailyBucket>::new();
|
||
|
|
for item in self
|
||
|
|
.by_request_id
|
||
|
|
.read()
|
||
|
|
.expect("usage repository lock")
|
||
|
|
.values()
|
||
|
|
{
|
||
|
|
if !usage_matches_dashboard_daily_breakdown_query(item, query) {
|
||
|
|
continue;
|
||
|
|
}
|
||
|
|
let Some(date) = usage_dashboard_local_date(item, query.tz_offset_minutes) else {
|
||
|
|
continue;
|
||
|
|
};
|
||
|
|
let key = (date, item.model.clone(), item.provider_name.clone());
|
||
|
|
let bucket = grouped.entry(key).or_default();
|
||
|
|
bucket.requests = bucket.requests.saturating_add(1);
|
||
|
|
bucket.total_tokens = bucket.total_tokens.saturating_add(item.total_tokens);
|
||
|
|
bucket.total_cost_usd += item.total_cost_usd;
|
||
|
|
if let Some(response_time_ms) = item.response_time_ms {
|
||
|
|
bucket.response_time_sum_ms += response_time_ms as f64;
|
||
|
|
bucket.response_time_samples = bucket.response_time_samples.saturating_add(1);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
Ok(grouped
|
||
|
|
.into_iter()
|
||
|
|
.map(
|
||
|
|
|((date, model, provider), bucket)| StoredUsageDashboardDailyBreakdownRow {
|
||
|
|
date,
|
||
|
|
model,
|
||
|
|
provider,
|
||
|
|
requests: bucket.requests,
|
||
|
|
total_tokens: bucket.total_tokens,
|
||
|
|
total_cost_usd: bucket.total_cost_usd,
|
||
|
|
response_time_sum_ms: bucket.response_time_sum_ms,
|
||
|
|
response_time_samples: bucket.response_time_samples,
|
||
|
|
},
|
||
|
|
)
|
||
|
|
.collect())
|
||
|
|
}
|
||
|
|
|
||
|
|
async fn summarize_dashboard_provider_counts(
|
||
|
|
&self,
|
||
|
|
query: &UsageDashboardProviderCountsQuery,
|
||
|
|
) -> Result<Vec<StoredUsageDashboardProviderCount>, DataLayerError> {
|
||
|
|
let mut grouped = BTreeMap::<String, u64>::new();
|
||
|
|
for item in self
|
||
|
|
.by_request_id
|
||
|
|
.read()
|
||
|
|
.expect("usage repository lock")
|
||
|
|
.values()
|
||
|
|
{
|
||
|
|
if !usage_matches_dashboard_provider_counts_query(item, query) {
|
||
|
|
continue;
|
||
|
|
}
|
||
|
|
*grouped.entry(item.provider_name.clone()).or_default() += 1;
|
||
|
|
}
|
||
|
|
|
||
|
|
let mut items = grouped
|
||
|
|
.into_iter()
|
||
|
|
.map(
|
||
|
|
|(provider_name, request_count)| StoredUsageDashboardProviderCount {
|
||
|
|
provider_name,
|
||
|
|
request_count,
|
||
|
|
},
|
||
|
|
)
|
||
|
|
.collect::<Vec<_>>();
|
||
|
|
items.sort_by(|left, right| {
|
||
|
|
right
|
||
|
|
.request_count
|
||
|
|
.cmp(&left.request_count)
|
||
|
|
.then_with(|| left.provider_name.cmp(&right.provider_name))
|
||
|
|
});
|
||
|
|
Ok(items)
|
||
|
|
}
|
||
|
|
|
||
|
|
async fn summarize_usage_breakdown(
|
||
|
|
&self,
|
||
|
|
query: &UsageBreakdownSummaryQuery,
|
||
|
|
) -> Result<Vec<StoredUsageBreakdownSummaryRow>, DataLayerError> {
|
||
|
|
#[derive(Default)]
|
||
|
|
struct BreakdownBucket {
|
||
|
|
request_count: u64,
|
||
|
|
input_tokens: u64,
|
||
|
|
total_tokens: u64,
|
||
|
|
output_tokens: u64,
|
||
|
|
effective_input_tokens: u64,
|
||
|
|
total_input_context: u64,
|
||
|
|
cache_creation_tokens: u64,
|
||
|
|
cache_creation_ephemeral_5m_tokens: u64,
|
||
|
|
cache_creation_ephemeral_1h_tokens: u64,
|
||
|
|
cache_read_tokens: u64,
|
||
|
|
total_cost_usd: f64,
|
||
|
|
actual_total_cost_usd: f64,
|
||
|
|
success_count: u64,
|
||
|
|
response_time_sum_ms: f64,
|
||
|
|
response_time_samples: u64,
|
||
|
|
overall_response_time_sum_ms: f64,
|
||
|
|
overall_response_time_samples: u64,
|
||
|
|
}
|
||
|
|
|
||
|
|
let mut grouped = BTreeMap::<String, BreakdownBucket>::new();
|
||
|
|
for item in self
|
||
|
|
.by_request_id
|
||
|
|
.read()
|
||
|
|
.expect("usage repository lock")
|
||
|
|
.values()
|
||
|
|
{
|
||
|
|
if !usage_matches_breakdown_summary_query(item, query) {
|
||
|
|
continue;
|
||
|
|
}
|
||
|
|
let group_key = match query.group_by {
|
||
|
|
UsageBreakdownGroupBy::Model => item.model.clone(),
|
||
|
|
UsageBreakdownGroupBy::Provider => item.provider_name.clone(),
|
||
|
|
UsageBreakdownGroupBy::ApiFormat => item.api_format.clone().unwrap_or_default(),
|
||
|
|
};
|
||
|
|
let bucket = grouped.entry(group_key).or_default();
|
||
|
|
let is_success = item.status != "failed"
|
||
|
|
&& item.status_code.is_none_or(|status| status < 400)
|
||
|
|
&& item.error_message.is_none();
|
||
|
|
|
||
|
|
bucket.request_count = bucket.request_count.saturating_add(1);
|
||
|
|
bucket.input_tokens = bucket.input_tokens.saturating_add(item.input_tokens);
|
||
|
|
bucket.total_tokens = bucket.total_tokens.saturating_add(item.total_tokens);
|
||
|
|
bucket.output_tokens = bucket.output_tokens.saturating_add(item.output_tokens);
|
||
|
|
bucket.effective_input_tokens = bucket
|
||
|
|
.effective_input_tokens
|
||
|
|
.saturating_add(usage_effective_input_tokens(item));
|
||
|
|
bucket.total_input_context = bucket
|
||
|
|
.total_input_context
|
||
|
|
.saturating_add(usage_total_input_context(item));
|
||
|
|
bucket.cache_creation_tokens = bucket
|
||
|
|
.cache_creation_tokens
|
||
|
|
.saturating_add(usage_cache_creation_tokens(item));
|
||
|
|
bucket.cache_creation_ephemeral_5m_tokens = bucket
|
||
|
|
.cache_creation_ephemeral_5m_tokens
|
||
|
|
.saturating_add(item.cache_creation_ephemeral_5m_input_tokens);
|
||
|
|
bucket.cache_creation_ephemeral_1h_tokens = bucket
|
||
|
|
.cache_creation_ephemeral_1h_tokens
|
||
|
|
.saturating_add(item.cache_creation_ephemeral_1h_input_tokens);
|
||
|
|
bucket.cache_read_tokens = bucket
|
||
|
|
.cache_read_tokens
|
||
|
|
.saturating_add(item.cache_read_input_tokens);
|
||
|
|
bucket.total_cost_usd += item.total_cost_usd;
|
||
|
|
bucket.actual_total_cost_usd += item.actual_total_cost_usd;
|
||
|
|
if let Some(response_time_ms) = item.response_time_ms {
|
||
|
|
bucket.overall_response_time_sum_ms += response_time_ms as f64;
|
||
|
|
bucket.overall_response_time_samples =
|
||
|
|
bucket.overall_response_time_samples.saturating_add(1);
|
||
|
|
}
|
||
|
|
if is_success {
|
||
|
|
bucket.success_count = bucket.success_count.saturating_add(1);
|
||
|
|
if let Some(response_time_ms) = item.response_time_ms {
|
||
|
|
bucket.response_time_sum_ms += response_time_ms as f64;
|
||
|
|
bucket.response_time_samples = bucket.response_time_samples.saturating_add(1);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
let mut items = grouped
|
||
|
|
.into_iter()
|
||
|
|
.map(|(group_key, bucket)| StoredUsageBreakdownSummaryRow {
|
||
|
|
group_key,
|
||
|
|
request_count: bucket.request_count,
|
||
|
|
input_tokens: bucket.input_tokens,
|
||
|
|
total_tokens: bucket.total_tokens,
|
||
|
|
output_tokens: bucket.output_tokens,
|
||
|
|
effective_input_tokens: bucket.effective_input_tokens,
|
||
|
|
total_input_context: bucket.total_input_context,
|
||
|
|
cache_creation_tokens: bucket.cache_creation_tokens,
|
||
|
|
cache_creation_ephemeral_5m_tokens: bucket.cache_creation_ephemeral_5m_tokens,
|
||
|
|
cache_creation_ephemeral_1h_tokens: bucket.cache_creation_ephemeral_1h_tokens,
|
||
|
|
cache_read_tokens: bucket.cache_read_tokens,
|
||
|
|
total_cost_usd: bucket.total_cost_usd,
|
||
|
|
actual_total_cost_usd: bucket.actual_total_cost_usd,
|
||
|
|
success_count: bucket.success_count,
|
||
|
|
response_time_sum_ms: bucket.response_time_sum_ms,
|
||
|
|
response_time_samples: bucket.response_time_samples,
|
||
|
|
overall_response_time_sum_ms: bucket.overall_response_time_sum_ms,
|
||
|
|
overall_response_time_samples: bucket.overall_response_time_samples,
|
||
|
|
})
|
||
|
|
.collect::<Vec<_>>();
|
||
|
|
items.sort_by(|left, right| {
|
||
|
|
right
|
||
|
|
.request_count
|
||
|
|
.cmp(&left.request_count)
|
||
|
|
.then_with(|| left.group_key.cmp(&right.group_key))
|
||
|
|
});
|
||
|
|
Ok(items)
|
||
|
|
}
|
||
|
|
|
||
|
|
async fn count_monitoring_usage_errors(
|
||
|
|
&self,
|
||
|
|
query: &UsageMonitoringErrorCountQuery,
|
||
|
|
) -> Result<u64, DataLayerError> {
|
||
|
|
Ok(self
|
||
|
|
.by_request_id
|
||
|
|
.read()
|
||
|
|
.expect("usage repository lock")
|
||
|
|
.values()
|
||
|
|
.filter(|item| usage_matches_monitoring_error_count_query(item, query))
|
||
|
|
.count() as u64)
|
||
|
|
}
|
||
|
|
|
||
|
|
async fn list_monitoring_usage_errors(
|
||
|
|
&self,
|
||
|
|
query: &UsageMonitoringErrorListQuery,
|
||
|
|
) -> Result<Vec<StoredRequestUsageAudit>, DataLayerError> {
|
||
|
|
let mut items: Vec<_> = self
|
||
|
|
.by_request_id
|
||
|
|
.read()
|
||
|
|
.expect("usage repository lock")
|
||
|
|
.values()
|
||
|
|
.filter(|item| usage_matches_monitoring_error_list_query(item, query))
|
||
|
|
.cloned()
|
||
|
|
.collect();
|
||
|
|
sort_usage_items(&mut items, true);
|
||
|
|
if let Some(limit) = query.limit {
|
||
|
|
items.truncate(limit);
|
||
|
|
}
|
||
|
|
Ok(items)
|
||
|
|
}
|
||
|
|
|
||
|
|
async fn summarize_usage_error_distribution(
|
||
|
|
&self,
|
||
|
|
query: &UsageErrorDistributionQuery,
|
||
|
|
) -> Result<Vec<StoredUsageErrorDistributionRow>, DataLayerError> {
|
||
|
|
let mut grouped = BTreeMap::<(String, String), u64>::new();
|
||
|
|
for item in self
|
||
|
|
.by_request_id
|
||
|
|
.read()
|
||
|
|
.expect("usage repository lock")
|
||
|
|
.values()
|
||
|
|
{
|
||
|
|
if !usage_matches_error_distribution_query(item, query) {
|
||
|
|
continue;
|
||
|
|
}
|
||
|
|
let Some(date) = usage_dashboard_local_date(item, query.tz_offset_minutes) else {
|
||
|
|
continue;
|
||
|
|
};
|
||
|
|
let Some(category) = item.error_category.clone() else {
|
||
|
|
continue;
|
||
|
|
};
|
||
|
|
*grouped.entry((date, category)).or_default() += 1;
|
||
|
|
}
|
||
|
|
|
||
|
|
Ok(grouped
|
||
|
|
.into_iter()
|
||
|
|
.map(
|
||
|
|
|((date, error_category), count)| StoredUsageErrorDistributionRow {
|
||
|
|
date,
|
||
|
|
error_category,
|
||
|
|
count,
|
||
|
|
},
|
||
|
|
)
|
||
|
|
.collect())
|
||
|
|
}
|
||
|
|
|
||
|
|
async fn summarize_usage_performance_percentiles(
|
||
|
|
&self,
|
||
|
|
query: &UsagePerformancePercentilesQuery,
|
||
|
|
) -> Result<Vec<StoredUsagePerformancePercentilesRow>, DataLayerError> {
|
||
|
|
let mut grouped = BTreeMap::<String, (Vec<u64>, Vec<u64>)>::new();
|
||
|
|
for item in self
|
||
|
|
.by_request_id
|
||
|
|
.read()
|
||
|
|
.expect("usage repository lock")
|
||
|
|
.values()
|
||
|
|
{
|
||
|
|
if !usage_matches_performance_percentiles_query(item, query) {
|
||
|
|
continue;
|
||
|
|
}
|
||
|
|
let Some(date) = usage_dashboard_local_date(item, query.tz_offset_minutes) else {
|
||
|
|
continue;
|
||
|
|
};
|
||
|
|
let entry = grouped.entry(date).or_default();
|
||
|
|
if let Some(response_time_ms) = item.response_time_ms {
|
||
|
|
entry.0.push(response_time_ms);
|
||
|
|
}
|
||
|
|
if let Some(first_byte_time_ms) = item.first_byte_time_ms {
|
||
|
|
entry.1.push(first_byte_time_ms);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
Ok(grouped
|
||
|
|
.into_iter()
|
||
|
|
.map(|(date, (mut response_times, mut first_byte_times))| {
|
||
|
|
StoredUsagePerformancePercentilesRow {
|
||
|
|
date,
|
||
|
|
p50_response_time_ms: usage_percentile_cont(&mut response_times, 0.5),
|
||
|
|
p90_response_time_ms: usage_percentile_cont(&mut response_times, 0.9),
|
||
|
|
p99_response_time_ms: usage_percentile_cont(&mut response_times, 0.99),
|
||
|
|
p50_first_byte_time_ms: usage_percentile_cont(&mut first_byte_times, 0.5),
|
||
|
|
p90_first_byte_time_ms: usage_percentile_cont(&mut first_byte_times, 0.9),
|
||
|
|
p99_first_byte_time_ms: usage_percentile_cont(&mut first_byte_times, 0.99),
|
||
|
|
}
|
||
|
|
})
|
||
|
|
.collect())
|
||
|
|
}
|
||
|
|
|
||
async fn summarize_usage_provider_performance(
|
|||
|
|
&self,
|
||
|
|
query: &UsageProviderPerformanceQuery,
|
||
|
|
) -> Result<StoredUsageProviderPerformance, DataLayerError> {
|
||
|
|
#[derive(Default)]
|
||
|
|
struct ProviderPerformanceBucket {
|
||
|
|
provider: String,
|
||
|
|
request_count: u64,
|
||
|
|
success_count: u64,
|
||
|
|
output_tokens: u64,
|
||
|
|
tps_output_tokens: u64,
|
||
|
|
tps_response_time_ms_sum: u64,
|
||
|
|
tps_sample_count: u64,
|
||
|
|
first_byte_time_ms_sum: u64,
|
||
|
|
first_byte_sample_count: u64,
|
||
|
|
response_time_ms_sum: u64,
|
||
|
|
response_time_sample_count: u64,
|
||
|
|
response_times: Vec<u64>,
|
||
|
|
first_byte_times: Vec<u64>,
|
||
slow_request_count: u64,
|
|||
}
|
|||
|
|
|
||
|
|
impl ProviderPerformanceBucket {
|
||
fn add(&mut self, item: &StoredRequestUsageAudit, slow_threshold_ms: u64) {
|
|||
self.request_count = self.request_count.saturating_add(1);
|
|||
|
|
self.output_tokens = self.output_tokens.saturating_add(item.output_tokens);
|
||
if item
|
|||
|
|
.response_time_ms
|
||
|
|
.is_some_and(|value| value >= slow_threshold_ms)
|
||
|
|
{
|
||
|
|
self.slow_request_count = self.slow_request_count.saturating_add(1);
|
||
|
|
}
|
||
if !usage_is_success(item) {
|
|||
|
|
return;
|
||
|
|
}
|
||
|
|
|
||
|
|
self.success_count = self.success_count.saturating_add(1);
|
||
|
|
if let Some(response_time_ms) = item.response_time_ms {
|
||
|
|
self.response_time_ms_sum =
|
||
|
|
self.response_time_ms_sum.saturating_add(response_time_ms);
|
||
|
|
self.response_time_sample_count =
|
||
|
|
self.response_time_sample_count.saturating_add(1);
|
||
|
|
self.response_times.push(response_time_ms);
|
||
if let Some(output_tps_duration_ms) = usage_output_tps_duration_ms(item) {
|
|||
|
|
if item.output_tokens > 0 {
|
||
|
|
self.tps_output_tokens =
|
||
|
|
self.tps_output_tokens.saturating_add(item.output_tokens);
|
||
|
|
self.tps_response_time_ms_sum = self
|
||
|
|
.tps_response_time_ms_sum
|
||
|
|
.saturating_add(output_tps_duration_ms);
|
||
|
|
self.tps_sample_count = self.tps_sample_count.saturating_add(1);
|
||
|
|
}
|
||
}
|
|||
|
|
}
|
||
|
|
if let Some(first_byte_time_ms) = item.first_byte_time_ms {
|
||
|
|
self.first_byte_time_ms_sum = self
|
||
|
|
.first_byte_time_ms_sum
|
||
|
|
.saturating_add(first_byte_time_ms);
|
||
|
|
self.first_byte_sample_count = self.first_byte_sample_count.saturating_add(1);
|
||
|
|
self.first_byte_times.push(first_byte_time_ms);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
fn avg(sum: u64, samples: u64) -> Option<f64> {
|
||
|
|
if samples == 0 {
|
||
|
|
None
|
||
|
|
} else {
|
||
|
|
Some(sum as f64 / samples as f64)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
fn avg_tps(tokens: u64, response_time_ms_sum: u64) -> Option<f64> {
|
||
|
|
if response_time_ms_sum == 0 {
|
||
|
|
None
|
||
|
|
} else {
|
||
|
|
Some(tokens as f64 * 1000.0 / response_time_ms_sum as f64)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
let usage = self
|
||
|
|
.by_request_id
|
||
|
|
.read()
|
||
|
|
.expect("usage repository lock")
|
||
|
|
.values()
|
||
|
|
.cloned()
|
||
|
|
.collect::<Vec<_>>();
|
||
|
|
|
||
|
|
let mut grouped = BTreeMap::<String, ProviderPerformanceBucket>::new();
|
||
|
|
let mut summary_bucket = ProviderPerformanceBucket::default();
|
||
|
|
for item in &usage {
|
||
|
|
let Some((provider_id, provider)) =
|
||
|
|
usage_matches_provider_performance_query(item, query)
|
||
|
|
else {
|
||
|
|
continue;
|
||
|
|
};
|
||
summary_bucket.add(item, query.slow_threshold_ms);
|
|||
let bucket = grouped.entry(provider_id).or_default();
|
|||
|
|
if bucket.provider.is_empty() {
|
||
|
|
bucket.provider = provider;
|
||
|
|
}
|
||
bucket.add(item, query.slow_threshold_ms);
|
|||
}
|
|||
|
|
|
||
let mut summary_response_times = summary_bucket.response_times.clone();
|
|||
|
|
let mut summary_first_byte_times = summary_bucket.first_byte_times.clone();
|
||
let summary = StoredUsageProviderPerformanceSummary {
|
|||
|
|
request_count: summary_bucket.request_count,
|
||
|
|
success_count: summary_bucket.success_count,
|
||
|
|
avg_output_tps: avg_tps(
|
||
|
|
summary_bucket.tps_output_tokens,
|
||
|
|
summary_bucket.tps_response_time_ms_sum,
|
||
|
|
),
|
||
|
|
avg_first_byte_time_ms: avg(
|
||
|
|
summary_bucket.first_byte_time_ms_sum,
|
||
|
|
summary_bucket.first_byte_sample_count,
|
||
|
|
),
|
||
|
|
avg_response_time_ms: avg(
|
||
|
|
summary_bucket.response_time_ms_sum,
|
||
|
|
summary_bucket.response_time_sample_count,
|
||
|
|
),
|
||
p90_response_time_ms: usage_percentile_cont(&mut summary_response_times, 0.9),
|
|||
|
|
p99_response_time_ms: usage_percentile_cont(&mut summary_response_times, 0.99),
|
||
|
|
p90_first_byte_time_ms: usage_percentile_cont(&mut summary_first_byte_times, 0.9),
|
||
|
|
p99_first_byte_time_ms: usage_percentile_cont(&mut summary_first_byte_times, 0.99),
|
||
|
|
tps_sample_count: summary_bucket.tps_sample_count,
|
||
|
|
response_time_sample_count: summary_bucket.response_time_sample_count,
|
||
|
|
first_byte_sample_count: summary_bucket.first_byte_sample_count,
|
||
|
|
slow_request_count: summary_bucket.slow_request_count,
|
||
};
|
|||
|
|
|
||
|
|
let mut providers = grouped
|
||
|
|
.into_iter()
|
||
|
|
.map(|(provider_id, mut bucket)| {
|
||
|
|
let p90_response_time_ms = usage_percentile_cont(&mut bucket.response_times, 0.9);
|
||
let p99_response_time_ms = usage_percentile_cont(&mut bucket.response_times, 0.99);
|
|||
let p90_first_byte_time_ms =
|
|||
|
|
usage_percentile_cont(&mut bucket.first_byte_times, 0.9);
|
||
let p99_first_byte_time_ms =
|
|||
|
|
usage_percentile_cont(&mut bucket.first_byte_times, 0.99);
|
||
StoredUsageProviderPerformanceProviderRow {
|
|||
|
|
provider_id,
|
||
|
|
provider: bucket.provider,
|
||
|
|
request_count: bucket.request_count,
|
||
|
|
success_count: bucket.success_count,
|
||
|
|
output_tokens: bucket.output_tokens,
|
||
|
|
avg_output_tps: avg_tps(
|
||
|
|
bucket.tps_output_tokens,
|
||
|
|
bucket.tps_response_time_ms_sum,
|
||
|
|
),
|
||
|
|
avg_first_byte_time_ms: avg(
|
||
|
|
bucket.first_byte_time_ms_sum,
|
||
|
|
bucket.first_byte_sample_count,
|
||
|
|
),
|
||
|
|
avg_response_time_ms: avg(
|
||
|
|
bucket.response_time_ms_sum,
|
||
|
|
bucket.response_time_sample_count,
|
||
|
|
),
|
||
|
|
p90_response_time_ms,
|
||
p99_response_time_ms,
|
|||
p90_first_byte_time_ms,
|
|||
p99_first_byte_time_ms,
|
|||
tps_sample_count: bucket.tps_sample_count,
|
|||
response_time_sample_count: bucket.response_time_sample_count,
|
|||
first_byte_sample_count: bucket.first_byte_sample_count,
|
|||
slow_request_count: bucket.slow_request_count,
|
|||
}
|
|||
|
|
})
|
||
|
|
.collect::<Vec<_>>();
|
||
|
|
providers.sort_by(|left, right| {
|
||
|
|
right
|
||
|
|
.request_count
|
||
|
|
.cmp(&left.request_count)
|
||
|
|
.then_with(|| left.provider_id.cmp(&right.provider_id))
|
||
|
|
});
|
||
|
|
providers.truncate(query.limit.max(1));
|
||
|
|
|
||
|
|
let top_provider_ids = providers
|
||
|
|
.iter()
|
||
|
|
.map(|row| row.provider_id.clone())
|
||
|
|
.collect::<Vec<_>>();
|
||
|
|
let mut timeline_grouped = BTreeMap::<(String, String), ProviderPerformanceBucket>::new();
|
||
|
|
for item in &usage {
|
||
|
|
let Some((provider_id, provider)) =
|
||
|
|
usage_matches_provider_performance_query(item, query)
|
||
|
|
else {
|
||
|
|
continue;
|
||
|
|
};
|
||
|
|
if !top_provider_ids.iter().any(|value| value == &provider_id) {
|
||
|
|
continue;
|
||
|
|
}
|
||
|
|
let Some(bucket_key) =
|
||
|
|
usage_time_series_bucket_key(item, query.granularity, query.tz_offset_minutes)
|
||
|
|
else {
|
||
|
|
continue;
|
||
|
|
};
|
||
|
|
let bucket = timeline_grouped
|
||
|
|
.entry((bucket_key, provider_id))
|
||
|
|
.or_default();
|
||
|
|
if bucket.provider.is_empty() {
|
||
|
|
bucket.provider = provider;
|
||
|
|
}
|
||
bucket.add(item, query.slow_threshold_ms);
|
|||
}
|
|||
|
|
|
||
|
|
let timeline = timeline_grouped
|
||
|
|
.into_iter()
|
||
|
|
.map(
|
||
|
|
|((date, provider_id), bucket)| StoredUsageProviderPerformanceTimelineRow {
|
||
|
|
date,
|
||
|
|
provider_id,
|
||
|
|
provider: bucket.provider,
|
||
|
|
request_count: bucket.request_count,
|
||
|
|
success_count: bucket.success_count,
|
||
|
|
output_tokens: bucket.output_tokens,
|
||
|
|
avg_output_tps: avg_tps(
|
||
|
|
bucket.tps_output_tokens,
|
||
|
|
bucket.tps_response_time_ms_sum,
|
||
|
|
),
|
||
|
|
avg_first_byte_time_ms: avg(
|
||
|
|
bucket.first_byte_time_ms_sum,
|
||
|
|
bucket.first_byte_sample_count,
|
||
|
|
),
|
||
|
|
avg_response_time_ms: avg(
|
||
|
|
bucket.response_time_ms_sum,
|
||
|
|
bucket.response_time_sample_count,
|
||
|
|
),
|
||
slow_request_count: bucket.slow_request_count,
|
|||
},
|
|||
|
|
)
|
||
|
|
.collect();
|
||
|
|
|
||
|
|
Ok(StoredUsageProviderPerformance {
|
||
|
|
summary,
|
||
|
|
providers,
|
||
|
|
timeline,
|
||
|
|
})
|
||
|
|
}
|
||
|
|
|
||
async fn summarize_usage_cost_savings(
|
|||
|
|
&self,
|
||
|
|
query: &UsageCostSavingsSummaryQuery,
|
||
|
|
) -> Result<StoredUsageCostSavingsSummary, DataLayerError> {
|
||
|
|
let mut summary = StoredUsageCostSavingsSummary::default();
|
||
|
|
for item in self
|
||
|
|
.by_request_id
|
||
|
|
.read()
|
||
|
|
.expect("usage repository lock")
|
||
|
|
.values()
|
||
|
|
{
|
||
|
|
if !usage_matches_cost_savings_query(item, query) {
|
||
|
|
continue;
|
||
|
|
}
|
||
|
|
summary.cache_read_tokens = summary
|
||
|
|
.cache_read_tokens
|
||
|
|
.saturating_add(item.cache_read_input_tokens);
|
||
|
|
summary.cache_read_cost_usd += item.cache_read_cost_usd;
|
||
|
|
summary.cache_creation_cost_usd += item.cache_creation_cost_usd;
|
||
summary.estimated_full_cost_usd += item.settlement_input_price_per_1m().unwrap_or(0.0)
|
|||
* item.cache_read_input_tokens as f64
|
|||
|
|
/ 1_000_000.0;
|
||
|
|
}
|
||
|
|
Ok(summary)
|
||
|
|
}
|
||
|
|
|
||
|
|
async fn summarize_usage_cache_affinity_hit_summary(
|
||
|
|
&self,
|
||
|
|
query: &UsageCacheAffinityHitSummaryQuery,
|
||
|
|
) -> Result<StoredUsageCacheAffinityHitSummary, DataLayerError> {
|
||
|
|
let mut summary = StoredUsageCacheAffinityHitSummary::default();
|
||
|
|
for item in self
|
||
|
|
.by_request_id
|
||
|
|
.read()
|
||
|
|
.expect("usage repository lock")
|
||
|
|
.values()
|
||
|
|
{
|
||
|
|
if !usage_matches_cache_affinity_hit_summary_query(item, query) {
|
||
|
|
continue;
|
||
|
|
}
|
||
|
|
summary.total_requests = summary.total_requests.saturating_add(1);
|
||
|
|
summary.input_tokens = summary.input_tokens.saturating_add(item.input_tokens);
|
||
|
|
summary.cache_read_tokens = summary
|
||
|
|
.cache_read_tokens
|
||
|
|
.saturating_add(item.cache_read_input_tokens);
|
||
|
|
summary.cache_creation_tokens = summary
|
||
|
|
.cache_creation_tokens
|
||
|
|
.saturating_add(usage_cache_creation_tokens(item));
|
||
|
|
summary.total_input_context = summary
|
||
|
|
.total_input_context
|
||
|
|
.saturating_add(usage_total_input_context(item));
|
||
|
|
summary.cache_read_cost_usd += item.cache_read_cost_usd;
|
||
|
|
summary.cache_creation_cost_usd += item.cache_creation_cost_usd;
|
||
|
|
if item.cache_read_input_tokens > 0 {
|
||
|
|
summary.requests_with_cache_hit = summary.requests_with_cache_hit.saturating_add(1);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
Ok(summary)
|
||
|
|
}
|
||
|
|
|
||
|
|
async fn list_usage_cache_affinity_intervals(
|
||
|
|
&self,
|
||
|
|
query: &UsageCacheAffinityIntervalQuery,
|
||
|
|
) -> Result<Vec<StoredUsageCacheAffinityIntervalRow>, DataLayerError> {
|
||
|
|
let mut items: Vec<_> = self
|
||
|
|
.by_request_id
|
||
|
|
.read()
|
||
|
|
.expect("usage repository lock")
|
||
|
|
.values()
|
||
|
|
.filter_map(|item| {
|
||
|
|
usage_matches_cache_affinity_interval_query(item, query)
|
||
|
|
.map(|group_id| (group_id, item.clone()))
|
||
|
|
})
|
||
|
|
.collect();
|
||
|
|
items.sort_by(|left, right| {
|
||
|
|
left.1
|
||
|
|
.created_at_unix_ms
|
||
|
|
.cmp(&right.1.created_at_unix_ms)
|
||
|
|
.then_with(|| left.1.id.cmp(&right.1.id))
|
||
|
|
});
|
||
|
|
|
||
|
|
let mut previous_created_at_by_group = BTreeMap::<String, u64>::new();
|
||
|
|
let mut rows = Vec::new();
|
||
|
|
for (group_id, item) in items {
|
||
|
|
let previous =
|
||
|
|
previous_created_at_by_group.insert(group_id.clone(), item.created_at_unix_ms);
|
||
|
|
let Some(previous_created_at_unix_secs) = previous else {
|
||
|
|
continue;
|
||
|
|
};
|
||
|
|
rows.push(StoredUsageCacheAffinityIntervalRow {
|
||
|
|
group_id,
|
||
|
|
username: item.username.clone(),
|
||
|
|
model: item.model.clone(),
|
||
|
|
created_at_unix_secs: item.created_at_unix_ms,
|
||
|
|
interval_minutes: item
|
||
|
|
.created_at_unix_ms
|
||
|
|
.saturating_sub(previous_created_at_unix_secs)
|
||
|
|
as f64
|
||
|
|
/ 60.0,
|
||
|
|
});
|
||
|
|
}
|
||
|
|
Ok(rows)
|
||
|
|
}
|
||
|
|
|
||
async fn summarize_usage_time_series(
|
|||
|
|
&self,
|
||
|
|
query: &UsageTimeSeriesQuery,
|
||
|
|
) -> Result<Vec<StoredUsageTimeSeriesBucket>, DataLayerError> {
|
||
|
|
let mut buckets = BTreeMap::<String, StoredUsageTimeSeriesBucket>::new();
|
||
|
|
for item in self
|
||
|
|
.by_request_id
|
||
|
|
.read()
|
||
|
|
.expect("usage repository lock")
|
||
|
|
.values()
|
||
|
|
{
|
||
|
|
if !usage_matches_time_series_query(item, query) {
|
||
|
|
continue;
|
||
|
|
}
|
||
|
|
let Some(bucket_key) =
|
||
|
|
usage_time_series_bucket_key(item, query.granularity, query.tz_offset_minutes)
|
||
|
|
else {
|
||
|
|
continue;
|
||
|
|
};
|
||
|
|
let bucket =
|
||
|
|
buckets
|
||
|
|
.entry(bucket_key.clone())
|
||
|
|
.or_insert_with(|| StoredUsageTimeSeriesBucket {
|
||
|
|
bucket_key,
|
||
|
|
..Default::default()
|
||
|
|
});
|
||
|
|
bucket.total_requests = bucket.total_requests.saturating_add(1);
|
||
|
|
bucket.input_tokens = bucket.input_tokens.saturating_add(item.input_tokens);
|
||
|
|
bucket.output_tokens = bucket.output_tokens.saturating_add(item.output_tokens);
|
||
|
|
bucket.cache_creation_tokens = bucket
|
||
|
|
.cache_creation_tokens
|
||
|
|
.saturating_add(item.cache_creation_input_tokens);
|
||
|
|
bucket.cache_read_tokens = bucket
|
||
|
|
.cache_read_tokens
|
||
|
|
.saturating_add(item.cache_read_input_tokens);
|
||
|
|
bucket.total_cost_usd += item.total_cost_usd;
|
||
|
|
bucket.total_response_time_ms += item.response_time_ms.unwrap_or(0) as f64;
|
||
|
|
}
|
||
|
|
Ok(buckets.into_values().collect())
|
||
|
|
}
|
||
|
|
|
||
|
|
async fn summarize_usage_leaderboard(
|
||
|
|
&self,
|
||
|
|
query: &UsageLeaderboardQuery,
|
||
|
|
) -> Result<Vec<StoredUsageLeaderboardSummary>, DataLayerError> {
|
||
|
|
let mut grouped = BTreeMap::<String, StoredUsageLeaderboardSummary>::new();
|
||
|
|
for item in self
|
||
|
|
.by_request_id
|
||
|
|
.read()
|
||
|
|
.expect("usage repository lock")
|
||
|
|
.values()
|
||
|
|
{
|
||
|
|
if !usage_matches_leaderboard_query(item, query) {
|
||
|
|
continue;
|
||
|
|
}
|
||
|
|
let (group_key, legacy_name) = match query.group_by {
|
||
|
|
UsageLeaderboardGroupBy::Model => (item.model.clone(), None),
|
||
|
|
UsageLeaderboardGroupBy::User => match item.user_id.clone() {
|
||
|
|
Some(user_id) => (user_id, item.username.clone()),
|
||
|
|
None => continue,
|
||
|
|
},
|
||
|
|
UsageLeaderboardGroupBy::ApiKey => match item.api_key_id.clone() {
|
||
|
|
Some(api_key_id) => (api_key_id, item.api_key_name.clone()),
|
||
|
|
None => continue,
|
||
|
|
},
|
||
|
|
};
|
||
|
|
let entry =
|
||
|
|
grouped
|
||
|
|
.entry(group_key.clone())
|
||
|
|
.or_insert_with(|| StoredUsageLeaderboardSummary {
|
||
|
|
group_key,
|
||
|
|
legacy_name: legacy_name.clone(),
|
||
|
|
..Default::default()
|
||
|
|
});
|
||
|
|
if entry.legacy_name.is_none() {
|
||
|
|
entry.legacy_name = legacy_name;
|
||
|
|
}
|
||
|
|
entry.request_count = entry.request_count.saturating_add(1);
|
||
entry.total_tokens = entry.total_tokens.saturating_add(usage_total_tokens(item));
|
|||
entry.total_cost_usd += item.total_cost_usd;
|
|||
|
|
}
|
||
|
|
Ok(grouped.into_values().collect())
|
||
|
|
}
|
||
|
|
|
||
async fn list_recent_usage_audits(
|
|||
|
|
&self,
|
||
|
|
user_id: Option<&str>,
|
||
|
|
limit: usize,
|
||
|
|
) -> Result<Vec<StoredRequestUsageAudit>, DataLayerError> {
|
||
|
|
let mut items: Vec<_> = self
|
||
|
|
.by_request_id
|
||
|
|
.read()
|
||
|
|
.expect("usage repository lock")
|
||
|
|
.values()
|
||
|
|
.filter(|item| match user_id {
|
||
|
|
Some(user_id) => item.user_id.as_deref() == Some(user_id),
|
||
|
|
None => true,
|
||
|
|
})
|
||
|
|
.cloned()
|
||
|
|
.collect();
|
||
|
|
items.sort_by(|left, right| {
|
||
|
|
right
|
||
.created_at_unix_ms
|
|||
|
|
.cmp(&left.created_at_unix_ms)
|
||
.then_with(|| left.id.cmp(&right.id))
|
|||
|
|
});
|
||
|
|
items.truncate(limit);
|
||
|
|
Ok(items)
|
||
|
|
}
|
||
|
|
|
||
|
|
async fn summarize_total_tokens_by_api_key_ids(
|
||
|
|
&self,
|
||
|
|
api_key_ids: &[String],
|
||
|
|
) -> Result<BTreeMap<String, u64>, DataLayerError> {
|
||
|
|
let api_key_id_set = api_key_ids.iter().map(String::as_str).collect::<Vec<_>>();
|
||
|
|
let mut totals = BTreeMap::<String, u64>::new();
|
||
|
|
for item in self
|
||
|
|
.by_request_id
|
||
|
|
.read()
|
||
|
|
.expect("usage repository lock")
|
||
|
|
.values()
|
||
|
|
{
|
||
|
|
let Some(api_key_id) = item.api_key_id.as_deref() else {
|
||
|
|
continue;
|
||
|
|
};
|
||
|
|
if !api_key_id_set.contains(&api_key_id) {
|
||
|
|
continue;
|
||
|
|
}
|
||
|
|
let entry = totals.entry(api_key_id.to_string()).or_insert(0);
|
||
|
|
*entry = (*entry).saturating_add(item.total_tokens);
|
||
|
|
}
|
||
|
|
Ok(totals)
|
||
|
|
}
|
||
|
|
|
||
async fn summarize_usage_totals_by_user_ids(
|
|||
|
|
&self,
|
||
|
|
user_ids: &[String],
|
||
|
|
) -> Result<Vec<StoredUsageUserTotals>, DataLayerError> {
|
||
|
|
let user_id_set = user_ids
|
||
|
|
.iter()
|
||
|
|
.cloned()
|
||
|
|
.collect::<std::collections::BTreeSet<_>>();
|
||
|
|
let mut totals = BTreeMap::<String, StoredUsageUserTotals>::new();
|
||
|
|
for item in self
|
||
|
|
.by_request_id
|
||
|
|
.read()
|
||
|
|
.expect("usage repository lock")
|
||
|
|
.values()
|
||
|
|
{
|
||
|
|
if matches!(item.status.as_str(), "pending" | "streaming")
|
||
|
|
|| matches!(item.provider_name.as_str(), "unknown" | "pending")
|
||
|
|
{
|
||
|
|
continue;
|
||
|
|
}
|
||
|
|
let Some(user_id) = item.user_id.as_deref() else {
|
||
|
|
continue;
|
||
|
|
};
|
||
|
|
if !user_id_set.contains(user_id) {
|
||
|
|
continue;
|
||
|
|
}
|
||
|
|
|
||
|
|
let entry =
|
||
|
|
totals
|
||
|
|
.entry(user_id.to_string())
|
||
|
|
.or_insert_with(|| StoredUsageUserTotals {
|
||
|
|
user_id: user_id.to_string(),
|
||
|
|
..Default::default()
|
||
|
|
});
|
||
|
|
entry.request_count = entry.request_count.saturating_add(1);
|
||
|
|
entry.total_tokens = entry.total_tokens.saturating_add(item.total_tokens);
|
||
|
|
}
|
||
|
|
Ok(totals.into_values().collect())
|
||
|
|
}
|
||
|
|
|
||
async fn summarize_usage_by_provider_api_key_ids(
|
|||
|
|
&self,
|
||
|
|
provider_api_key_ids: &[String],
|
||
|
|
) -> Result<BTreeMap<String, StoredProviderApiKeyUsageSummary>, DataLayerError> {
|
||
|
|
let provider_api_key_id_set = provider_api_key_ids
|
||
|
|
.iter()
|
||
|
|
.map(String::as_str)
|
||
|
|
.collect::<Vec<_>>();
|
||
|
|
let mut summaries = BTreeMap::<String, StoredProviderApiKeyUsageSummary>::new();
|
||
|
|
for item in self
|
||
|
|
.by_request_id
|
||
|
|
.read()
|
||
|
|
.expect("usage repository lock")
|
||
|
|
.values()
|
||
|
|
{
|
||
|
|
let Some(provider_api_key_id) = item.provider_api_key_id.as_deref() else {
|
||
|
|
continue;
|
||
|
|
};
|
||
|
|
if !provider_api_key_id_set.contains(&provider_api_key_id) {
|
||
|
|
continue;
|
||
|
|
}
|
||
|
|
let entry = summaries
|
||
|
|
.entry(provider_api_key_id.to_string())
|
||
|
|
.or_insert_with(|| StoredProviderApiKeyUsageSummary {
|
||
|
|
provider_api_key_id: provider_api_key_id.to_string(),
|
||
|
|
..StoredProviderApiKeyUsageSummary::default()
|
||
|
|
});
|
||
|
|
entry.request_count = entry.request_count.saturating_add(1);
|
||
|
|
entry.total_tokens = entry.total_tokens.saturating_add(item.total_tokens);
|
||
|
|
entry.total_cost_usd += item.total_cost_usd;
|
||
|
|
entry.last_used_at_unix_secs = Some(
|
||
|
|
entry
|
||
|
|
.last_used_at_unix_secs
|
||
|
|
.unwrap_or_default()
|
||
|
|
.max(item.created_at_unix_ms),
|
||
|
|
);
|
||
|
|
}
|
||
|
|
Ok(summaries)
|
||
|
|
}
|
||
|
|
|
||
async fn summarize_usage_by_provider_api_key_windows(
|
|||
|
|
&self,
|
||
|
|
requests: &[ProviderApiKeyWindowUsageRequest],
|
||
|
|
) -> Result<Vec<StoredProviderApiKeyWindowUsageSummary>, DataLayerError> {
|
||
|
|
let usage = self.by_request_id.read().expect("usage repository lock");
|
||
|
|
let mut summaries = Vec::with_capacity(requests.len());
|
||
|
|
|
||
|
|
for request in requests {
|
||
|
|
let provider_api_key_id = request.provider_api_key_id.trim();
|
||
|
|
if provider_api_key_id.is_empty() {
|
||
|
|
return Err(DataLayerError::InvalidInput(
|
||
|
|
"provider api key window usage provider_api_key_id cannot be empty".to_string(),
|
||
|
|
));
|
||
|
|
}
|
||
|
|
let window_code = request.window_code.trim();
|
||
|
|
if window_code.is_empty() {
|
||
|
|
return Err(DataLayerError::InvalidInput(
|
||
|
|
"provider api key window usage window_code cannot be empty".to_string(),
|
||
|
|
));
|
||
|
|
}
|
||
|
|
if request.start_unix_secs >= request.end_unix_secs {
|
||
|
|
return Err(DataLayerError::InvalidInput(
|
||
|
|
"provider api key window usage range must be non-empty".to_string(),
|
||
|
|
));
|
||
|
|
}
|
||
|
|
|
||
|
|
let mut summary = StoredProviderApiKeyWindowUsageSummary {
|
||
|
|
provider_api_key_id: provider_api_key_id.to_string(),
|
||
|
|
window_code: window_code.to_string(),
|
||
|
|
..StoredProviderApiKeyWindowUsageSummary::default()
|
||
|
|
};
|
||
|
|
|
||
|
|
for item in usage.values() {
|
||
|
|
if item.provider_api_key_id.as_deref() != Some(provider_api_key_id) {
|
||
|
|
continue;
|
||
|
|
}
|
||
|
|
if item.created_at_unix_ms < request.start_unix_secs
|
||
|
|
|| item.created_at_unix_ms >= request.end_unix_secs
|
||
|
|
{
|
||
|
|
continue;
|
||
|
|
}
|
||
|
|
|
||
|
|
summary.request_count = summary.request_count.saturating_add(1);
|
||
|
|
summary.total_tokens = summary.total_tokens.saturating_add(item.total_tokens);
|
||
|
|
summary.total_cost_usd += item.total_cost_usd;
|
||
|
|
}
|
||
|
|
|
||
|
|
summaries.push(summary);
|
||
|
|
}
|
||
|
|
|
||
|
|
Ok(summaries)
|
||
|
|
}
|
||
|
|
|
||
async fn summarize_provider_usage_since(
|
|||
|
|
&self,
|
||
|
|
provider_id: &str,
|
||
|
|
since_unix_secs: u64,
|
||
|
|
) -> Result<StoredProviderUsageSummary, DataLayerError> {
|
||
|
|
let windows = self
|
||
|
|
.provider_usage_windows
|
||
|
|
.read()
|
||
|
|
.expect("provider usage repository lock");
|
||
|
|
|
||
|
|
let mut summary = StoredProviderUsageSummary::default();
|
||
|
|
let mut response_time_samples = 0u64;
|
||
|
|
for window in windows.iter().filter(|window| {
|
||
|
|
window.provider_id == provider_id && window.window_start_unix_secs >= since_unix_secs
|
||
|
|
}) {
|
||
|
|
summary.total_requests = summary.total_requests.saturating_add(window.total_requests);
|
||
|
|
summary.successful_requests = summary
|
||
|
|
.successful_requests
|
||
|
|
.saturating_add(window.successful_requests);
|
||
|
|
summary.failed_requests = summary
|
||
|
|
.failed_requests
|
||
|
|
.saturating_add(window.failed_requests);
|
||
|
|
summary.total_cost_usd += window.total_cost_usd;
|
||
|
|
summary.avg_response_time_ms += window.avg_response_time_ms;
|
||
|
|
response_time_samples = response_time_samples.saturating_add(1);
|
||
|
|
}
|
||
|
|
|
||
|
|
if response_time_samples > 0 {
|
||
|
|
summary.avg_response_time_ms /= response_time_samples as f64;
|
||
|
|
}
|
||
|
|
|
||
|
|
Ok(summary)
|
||
|
|
}
|
||
|
|||
|
|
async fn summarize_usage_daily_heatmap(
|
||
|
|
&self,
|
||
|
|
query: &UsageDailyHeatmapQuery,
|
||
|
|
) -> Result<Vec<StoredUsageDailySummary>, DataLayerError> {
|
||
|
|
let items = self.by_request_id.read().expect("usage repository lock");
|
||
|
|
let mut daily = BTreeMap::<String, (u64, u64, f64, f64)>::new();
|
||
|
|
for item in items.values() {
|
||
|
|
if item.created_at_unix_ms < query.created_from_unix_secs {
|
||
|
|
continue;
|
||
|
|
}
|
||
|
|
if let Some(user_id) = &query.user_id {
|
||
|
|
if item.user_id.as_deref() != Some(user_id) {
|
||
|
|
continue;
|
||
|
|
}
|
||
|
|
}
|
||
if item.status == "pending" || item.status == "streaming" {
|
|||
|
|
continue;
|
||
|
|
}
|
||
|
|
if item.provider_name.eq_ignore_ascii_case("unknown")
|
||
|
|
|| item.provider_name.eq_ignore_ascii_case("pending")
|
||
|
|
{
|
||
continue;
|
|||
|
|
}
|
||
|
|
let ts = i64::try_from(item.created_at_unix_ms).unwrap_or_default();
|
||
|
|
let Some(dt) = chrono::DateTime::<chrono::Utc>::from_timestamp(ts, 0) else {
|
||
|
|
continue;
|
||
|
|
};
|
||
|
|
let date_key = dt.date_naive().to_string();
|
||
|
|
let entry = daily.entry(date_key).or_insert((0, 0, 0.0, 0.0));
|
||
|
|
entry.0 += 1;
|
||
|
|
let cache_creation = if item.cache_creation_input_tokens == 0
|
||
|
|
&& (item.cache_creation_ephemeral_5m_input_tokens
|
||
|
|
+ item.cache_creation_ephemeral_1h_input_tokens)
|
||
|
|
> 0
|
||
|
|
{
|
||
|
|
item.cache_creation_ephemeral_5m_input_tokens
|
||
|
|
+ item.cache_creation_ephemeral_1h_input_tokens
|
||
|
|
} else {
|
||
|
|
item.cache_creation_input_tokens
|
||
|
|
};
|
||
|
|
entry.1 += item.input_tokens
|
||
|
|
+ item.output_tokens
|
||
|
|
+ cache_creation
|
||
|
|
+ item.cache_read_input_tokens;
|
||
|
|
entry.2 += item.total_cost_usd;
|
||
|
|
entry.3 += item.actual_total_cost_usd;
|
||
|
|
}
|
||
|
|
let mut result: Vec<_> = daily
|
||
|
|
.into_iter()
|
||
|
|
.map(
|
||
|
|
|(date, (requests, total_tokens, total_cost_usd, actual_total_cost_usd))| {
|
||
|
|
StoredUsageDailySummary {
|
||
|
|
date,
|
||
|
|
requests,
|
||
|
|
total_tokens,
|
||
|
|
total_cost_usd,
|
||
|
|
actual_total_cost_usd,
|
||
|
|
}
|
||
|
|
},
|
||
|
|
)
|
||
|
|
.collect();
|
||
|
|
result.sort_by(|a, b| a.date.cmp(&b.date));
|
||
|
|
Ok(result)
|
||
|
|
}
|
||
}
|
|||
|
|
|
||
fn detach_usage_body(
|
|||
|
|
request_id: &str,
|
||
|
|
body: &mut Option<Value>,
|
||
|
|
detached_bodies: &mut BTreeMap<String, Value>,
|
||
|
|
field: UsageBodyField,
|
||
|
|
) -> Option<String> {
|
||
|
|
let value = body.take()?;
|
||
|
|
let body_ref = usage_body_ref(request_id, field);
|
||
|
|
detached_bodies.insert(body_ref.clone(), value);
|
||
|
|
Some(body_ref)
|
||
|
|
}
|
||
|
|
|
||
|
|
fn usage_body_ref_from_metadata(
|
||
|
|
metadata: Option<&Value>,
|
||
|
|
request_id: &str,
|
||
|
|
field: UsageBodyField,
|
||
|
|
) -> Option<String> {
|
||
|
|
metadata
|
||
|
|
.and_then(Value::as_object)
|
||
|
|
.and_then(|object| object.get(field.as_ref_key()))
|
||
|
|
.and_then(Value::as_str)
|
||
|
|
.map(str::trim)
|
||
|
|
.filter(|value| !value.is_empty())
|
||
|
|
.and_then(parse_usage_body_ref)
|
||
|
|
.filter(|(parsed_request_id, parsed_field)| {
|
||
|
|
parsed_request_id == request_id && *parsed_field == field
|
||
|
|
})
|
||
|
|
.map(|(parsed_request_id, parsed_field)| usage_body_ref(&parsed_request_id, parsed_field))
|
||
|
|
}
|
||
|
|
|
||
|
|
fn hydrate_legacy_body_refs(item: &mut StoredRequestUsageAudit) {
|
||
|
|
if item.request_body_ref.is_none() {
|
||
|
|
item.request_body_ref = usage_body_ref_from_metadata(
|
||
|
|
item.request_metadata.as_ref(),
|
||
|
|
&item.request_id,
|
||
|
|
UsageBodyField::RequestBody,
|
||
|
|
);
|
||
|
|
}
|
||
|
|
if item.provider_request_body_ref.is_none() {
|
||
|
|
item.provider_request_body_ref = usage_body_ref_from_metadata(
|
||
|
|
item.request_metadata.as_ref(),
|
||
|
|
&item.request_id,
|
||
|
|
UsageBodyField::ProviderRequestBody,
|
||
|
|
);
|
||
|
|
}
|
||
|
|
if item.response_body_ref.is_none() {
|
||
|
|
item.response_body_ref = usage_body_ref_from_metadata(
|
||
|
|
item.request_metadata.as_ref(),
|
||
|
|
&item.request_id,
|
||
|
|
UsageBodyField::ResponseBody,
|
||
|
|
);
|
||
|
|
}
|
||
|
|
if item.client_response_body_ref.is_none() {
|
||
|
|
item.client_response_body_ref = usage_body_ref_from_metadata(
|
||
|
|
item.request_metadata.as_ref(),
|
||
|
|
&item.request_id,
|
||
|
|
UsageBodyField::ClientResponseBody,
|
||
|
|
);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
fn hydrate_client_family(item: &mut StoredRequestUsageAudit) {
|
|||
|
|
if item.client_family.is_none() {
|
||
|
|
item.client_family = usage_request_metadata_client_family(item.request_metadata.as_ref())
|
||
|
|
.map(ToOwned::to_owned);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
fn persisted_usage_body_ref(
|
|||
|
|
incoming_ref: Option<&str>,
|
||
|
|
incoming_body: Option<&Value>,
|
||
|
|
_metadata: Option<&Value>,
|
||
|
|
existing: Option<&StoredRequestUsageAudit>,
|
||
|
|
field: UsageBodyField,
|
||
|
|
) -> Option<String> {
|
||
|
|
if incoming_body.is_some() {
|
||
|
|
return None;
|
||
|
|
}
|
||
|
|
incoming_ref
|
||
|
|
.map(str::trim)
|
||
|
|
.filter(|value| !value.is_empty())
|
||
|
|
.map(ToOwned::to_owned)
|
||
|
|
.or_else(|| {
|
||
|
|
existing.and_then(|existing| match field {
|
||
|
|
UsageBodyField::RequestBody => existing.request_body_ref.clone(),
|
||
|
|
UsageBodyField::ProviderRequestBody => existing.provider_request_body_ref.clone(),
|
||
|
|
UsageBodyField::ResponseBody => existing.response_body_ref.clone(),
|
||
|
|
UsageBodyField::ClientResponseBody => existing.client_response_body_ref.clone(),
|
||
|
|
})
|
||
|
|
})
|
||
|
|
}
|
||
|
|
|
||
fn merge_usage_status_code(
|
|||
|
|
existing: Option<&StoredRequestUsageAudit>,
|
||
|
|
incoming_status: &str,
|
||
|
|
incoming_status_code: Option<u16>,
|
||
|
|
) -> Option<u16> {
|
||
|
|
if existing.is_some_and(|existing| {
|
||
|
|
existing.status == "streaming"
|
||
|
|
&& incoming_status == "streaming"
|
||
|
|
&& incoming_status_code.is_none()
|
||
|
|
}) {
|
||
|
|
return existing.and_then(|existing| existing.status_code);
|
||
|
|
}
|
||
|
|
incoming_status_code
|
||
|
|
}
|
||
|
|
|
||
#[async_trait]
|
|||
|
|
impl UsageWriteRepository for InMemoryUsageReadRepository {
|
||
|
|
async fn upsert(
|
||
|
|
&self,
|
||
|
|
usage: UpsertUsageRecord,
|
||
|
|
) -> Result<StoredRequestUsageAudit, DataLayerError> {
|
||
|
|
usage.validate()?;
|
||
let usage = strip_deprecated_usage_display_fields(usage);
|
|||
let mut by_request_id = self.by_request_id.write().expect("usage repository lock");
|
|||
let existing = by_request_id.get(&usage.request_id).cloned();
|
|||
|
|||
let created_at_unix_ms = by_request_id
|
|||
.get(&usage.request_id)
|
|||
.map(|existing| existing.created_at_unix_ms)
|
|||
|
|
.or(usage.created_at_unix_ms)
|
||
.unwrap_or(usage.updated_at_unix_secs);
|
|||
|
|
|
||
|
|
let total_tokens = usage
|
||
|
|
.total_tokens
|
||
|
|
.or_else(|| {
|
||
|
|
Some(
|
||
|
|
usage.input_tokens.unwrap_or_default()
|
||
|
|
+ usage.output_tokens.unwrap_or_default(),
|
||
|
|
)
|
||
|
|
})
|
||
|
|
.unwrap_or_default();
|
||
if existing.as_ref().is_some_and(|existing| {
|
|||
usage_status_is_finalized(existing.status.as_str())
|
|||
|
|
&& usage_status_is_lifecycle(usage.status.as_str())
|
||
&& !usage_can_recover_terminal_failure(
|
|||
|
|
existing.status.as_str(),
|
||
|
|
existing.billing_status.as_str(),
|
||
|
|
usage.status.as_str(),
|
||
|
|
usage.billing_status.as_str(),
|
||
|
|
)
|
||
}) {
|
|||
|
|
return Ok(existing.expect("existing usage should be present").clone());
|
||
|
|
}
|
||
if existing.as_ref().is_some_and(|existing| {
|
|||
existing.billing_status == "pending"
|
|||
|
|
&& existing.status == "streaming"
|
||
|
|
&& usage.status == "pending"
|
||
|
|
}) {
|
||
|
|
return Ok(existing.expect("existing usage should be present").clone());
|
||
|
|
}
|
||
|
|||
let request_metadata = usage.request_metadata.clone().or_else(|| {
|
|||
|
|
existing
|
||
|
|
.as_ref()
|
||
|
|
.and_then(|existing| existing.request_metadata.clone())
|
||
|
|
});
|
||
let request_body_ref = persisted_usage_body_ref(
|
|||
|
|
usage.request_body_ref.as_deref(),
|
||
|
|
usage.request_body.as_ref(),
|
||
|
|
request_metadata.as_ref(),
|
||
existing.as_ref(),
|
|||
UsageBodyField::RequestBody,
|
|||
|
|
);
|
||
|
|
let provider_request_body_ref = persisted_usage_body_ref(
|
||
|
|
usage.provider_request_body_ref.as_deref(),
|
||
|
|
usage.provider_request_body.as_ref(),
|
||
|
|
request_metadata.as_ref(),
|
||
existing.as_ref(),
|
|||
UsageBodyField::ProviderRequestBody,
|
|||
|
|
);
|
||
|
|
let response_body_ref = persisted_usage_body_ref(
|
||
|
|
usage.response_body_ref.as_deref(),
|
||
|
|
usage.response_body.as_ref(),
|
||
|
|
request_metadata.as_ref(),
|
||
existing.as_ref(),
|
|||
UsageBodyField::ResponseBody,
|
|||
|
|
);
|
||
|
|
let client_response_body_ref = persisted_usage_body_ref(
|
||
|
|
usage.client_response_body_ref.as_deref(),
|
||
|
|
usage.client_response_body.as_ref(),
|
||
|
|
request_metadata.as_ref(),
|
||
existing.as_ref(),
|
|||
UsageBodyField::ClientResponseBody,
|
|||
|
|
);
|
||
let stored = StoredRequestUsageAudit {
|
|||
|
|
id: existing
|
||
.as_ref()
|
|||
.map(|existing| existing.id.clone())
|
|||
|
|
.unwrap_or_else(|| format!("usage-{}", usage.request_id)),
|
||
|
|
request_id: usage.request_id.clone(),
|
||
|
|
user_id: usage.user_id,
|
||
|
|
api_key_id: usage.api_key_id,
|
||
username: existing
|
|||
|
|
.as_ref()
|
||
|
|
.and_then(|existing| existing.username.clone()),
|
||
|
|
api_key_name: existing
|
||
|
|
.as_ref()
|
||
|
|
.and_then(|existing| existing.api_key_name.clone()),
|
||
provider_name: usage.provider_name,
|
|||
|
|
model: usage.model,
|
||
|
|
target_model: usage.target_model,
|
||
|
|
provider_id: usage.provider_id,
|
||
|
|
provider_endpoint_id: usage.provider_endpoint_id,
|
||
|
|
provider_api_key_id: usage.provider_api_key_id,
|
||
|
|
request_type: usage.request_type,
|
||
|
|
api_format: usage.api_format,
|
||
|
|
api_family: usage.api_family,
|
||
|
|
endpoint_kind: usage.endpoint_kind,
|
||
|
|
endpoint_api_format: usage.endpoint_api_format,
|
||
|
|
provider_api_family: usage.provider_api_family,
|
||
|
|
provider_endpoint_kind: usage.provider_endpoint_kind,
|
||
|
|
has_format_conversion: usage.has_format_conversion.unwrap_or(false),
|
||
|
|
is_stream: usage.is_stream.unwrap_or(false),
|
||
|
|
input_tokens: usage.input_tokens.unwrap_or_default(),
|
||
|
|
output_tokens: usage.output_tokens.unwrap_or_default(),
|
||
|
|
total_tokens,
|
||
|
|
cache_creation_input_tokens: usage.cache_creation_input_tokens.unwrap_or_else(|| {
|
||
|
|
existing
|
||
.as_ref()
|
|||
.map(|existing| existing.cache_creation_input_tokens)
|
|||
|
|
.unwrap_or_default()
|
||
|
|
}),
|
||
cache_creation_ephemeral_5m_input_tokens: usage
|
|||
|
|
.cache_creation_ephemeral_5m_input_tokens
|
||
|
|
.unwrap_or_else(|| {
|
||
|
|
existing
|
||
.as_ref()
|
|||
.map(|existing| existing.cache_creation_ephemeral_5m_input_tokens)
|
|||
|
|
.unwrap_or_default()
|
||
|
|
}),
|
||
|
|
cache_creation_ephemeral_1h_input_tokens: usage
|
||
|
|
.cache_creation_ephemeral_1h_input_tokens
|
||
|
|
.unwrap_or_else(|| {
|
||
|
|
existing
|
||
.as_ref()
|
|||
.map(|existing| existing.cache_creation_ephemeral_1h_input_tokens)
|
|||
|
|
.unwrap_or_default()
|
||
|
|
}),
|
||
cache_read_input_tokens: usage.cache_read_input_tokens.unwrap_or_else(|| {
|
|||
|
|
existing
|
||
.as_ref()
|
|||
.map(|existing| existing.cache_read_input_tokens)
|
|||
|
|
.unwrap_or_default()
|
||
|
|
}),
|
||
|
|
cache_creation_cost_usd: usage.cache_creation_cost_usd.unwrap_or_else(|| {
|
||
|
|
existing
|
||
.as_ref()
|
|||
.map(|existing| existing.cache_creation_cost_usd)
|
|||
|
|
.unwrap_or_default()
|
||
|
|
}),
|
||
|
|
cache_read_cost_usd: usage.cache_read_cost_usd.unwrap_or_else(|| {
|
||
|
|
existing
|
||
.as_ref()
|
|||
.map(|existing| existing.cache_read_cost_usd)
|
|||
|
|
.unwrap_or_default()
|
||
|
|
}),
|
||
output_price_per_1m: existing
|
|||
|
|
.as_ref()
|
||
|
|
.and_then(|existing| existing.output_price_per_1m),
|
||
total_cost_usd: usage.total_cost_usd.unwrap_or_else(|| {
|
|||
|
|
existing
|
||
.as_ref()
|
|||
.map(|existing| existing.total_cost_usd)
|
|||
|
|
.unwrap_or_default()
|
||
|
|
}),
|
||
|
|
actual_total_cost_usd: usage.actual_total_cost_usd.unwrap_or_else(|| {
|
||
|
|
existing
|
||
.as_ref()
|
|||
.map(|existing| existing.actual_total_cost_usd)
|
|||
|
|
.unwrap_or_default()
|
||
|
|
}),
|
||
status_code: merge_usage_status_code(
|
|||
|
|
existing.as_ref(),
|
||
|
|
usage.status.as_str(),
|
||
|
|
usage.status_code,
|
||
|
|
),
|
||
error_message: usage.error_message,
|
|||
|
|
error_category: usage.error_category,
|
||
response_time_ms: merge_usage_timing(
|
|||
|
|
existing
|
||
|
|
.as_ref()
|
||
|
|
.and_then(|existing| existing.response_time_ms),
|
||
|
|
usage.response_time_ms,
|
||
|
|
),
|
||
|
|
first_byte_time_ms: merge_usage_timing(
|
||
|
|
existing
|
||
|
|
.as_ref()
|
||
|
|
.and_then(|existing| existing.first_byte_time_ms),
|
||
|
|
usage.first_byte_time_ms,
|
||
|
|
),
|
||
status: usage.status,
|
|||
|
|
billing_status: usage.billing_status,
|
||
request_headers: usage.request_headers.or_else(|| {
|
|||
|
|
existing
|
||
|
|
.as_ref()
|
||
|
|
.and_then(|existing| existing.request_headers.clone())
|
||
|
|
}),
|
||
|
|
request_body: usage.request_body.or_else(|| {
|
||
|
|
existing
|
||
|
|
.as_ref()
|
||
|
|
.and_then(|existing| existing.request_body.clone())
|
||
|
|
}),
|
||
request_body_ref,
|
|||
request_body_state: usage.request_body_state.or_else(|| {
|
|||
|
|
existing
|
||
|
|
.as_ref()
|
||
|
|
.and_then(|existing| existing.request_body_state)
|
||
|
|
}),
|
||
provider_request_headers: usage.provider_request_headers.or_else(|| {
|
|||
existing
|
|||
|
|
.as_ref()
|
||
|
|
.and_then(|existing| existing.provider_request_headers.clone())
|
||
|
|
}),
|
||
|
|
provider_request_body: usage.provider_request_body.or_else(|| {
|
||
|
|
existing
|
||
|
|
.as_ref()
|
||
|
|
.and_then(|existing| existing.provider_request_body.clone())
|
||
}),
|
|||
provider_request_body_ref,
|
|||
provider_request_body_state: usage.provider_request_body_state.or_else(|| {
|
|||
|
|
existing
|
||
|
|
.as_ref()
|
||
|
|
.and_then(|existing| existing.provider_request_body_state)
|
||
|
|
}),
|
||
|
|
response_headers: usage.response_headers.or_else(|| {
|
||
|
|
existing
|
||
|
|
.as_ref()
|
||
|
|
.and_then(|existing| existing.response_headers.clone())
|
||
|
|
}),
|
||
|
|
response_body: usage.response_body.or_else(|| {
|
||
|
|
existing
|
||
|
|
.as_ref()
|
||
|
|
.and_then(|existing| existing.response_body.clone())
|
||
|
|
}),
|
||
response_body_ref,
|
|||
response_body_state: usage.response_body_state.or_else(|| {
|
|||
|
|
existing
|
||
|
|
.as_ref()
|
||
|
|
.and_then(|existing| existing.response_body_state)
|
||
|
|
}),
|
||
|
|
client_response_headers: usage.client_response_headers.or_else(|| {
|
||
|
|
existing
|
||
|
|
.as_ref()
|
||
|
|
.and_then(|existing| existing.client_response_headers.clone())
|
||
|
|
}),
|
||
|
|
client_response_body: usage.client_response_body.or_else(|| {
|
||
|
|
existing
|
||
|
|
.as_ref()
|
||
|
|
.and_then(|existing| existing.client_response_body.clone())
|
||
|
|
}),
|
||
client_response_body_ref,
|
|||
client_response_body_state: usage.client_response_body_state.or_else(|| {
|
|||
|
|
existing
|
||
|
|
.as_ref()
|
||
|
|
.and_then(|existing| existing.client_response_body_state)
|
||
|
|
}),
|
||
candidate_id: usage.candidate_id.or_else(|| {
|
|||
existing
|
|||
|
|
.as_ref()
|
||
|
|
.and_then(|existing| existing.routing_candidate_id().map(ToOwned::to_owned))
|
||
|
|
}),
|
||
|
|
candidate_index: usage.candidate_index.or_else(|| {
|
||
|
|
existing
|
||
|
|
.as_ref()
|
||
|
|
.and_then(|existing| existing.routing_candidate_index())
|
||
}),
|
|||
|
|
key_name: usage.key_name.or_else(|| {
|
||
existing
|
|||
|
|
.as_ref()
|
||
|
|
.and_then(|existing| existing.routing_key_name().map(ToOwned::to_owned))
|
||
}),
|
|||
|
|
planner_kind: usage.planner_kind.or_else(|| {
|
||
existing
|
|||
|
|
.as_ref()
|
||
|
|
.and_then(|existing| existing.routing_planner_kind().map(ToOwned::to_owned))
|
||
}),
|
|||
|
|
route_family: usage.route_family.or_else(|| {
|
||
existing
|
|||
|
|
.as_ref()
|
||
|
|
.and_then(|existing| existing.routing_route_family().map(ToOwned::to_owned))
|
||
}),
|
|||
|
|
route_kind: usage.route_kind.or_else(|| {
|
||
existing
|
|||
|
|
.as_ref()
|
||
|
|
.and_then(|existing| existing.routing_route_kind().map(ToOwned::to_owned))
|
||
}),
|
|||
|
|
execution_path: usage.execution_path.or_else(|| {
|
||
|
|
existing
|
||
.as_ref()
|
|||
.and_then(|existing| existing.routing_execution_path().map(ToOwned::to_owned))
|
|||
|
|
}),
|
||
|
|
local_execution_runtime_miss_reason: usage.local_execution_runtime_miss_reason.or_else(
|
||
|
|
|| {
|
||
existing.as_ref().and_then(|existing| {
|
|||
existing
|
|||
|
|
.routing_local_execution_runtime_miss_reason()
|
||
|
|
.map(ToOwned::to_owned)
|
||
|
|
})
|
||
|
|
},
|
||
|
|
),
|
||
client_family: usage_request_metadata_client_family(request_metadata.as_ref())
|
|||
|
|
.map(ToOwned::to_owned)
|
||
|
|
.or_else(|| {
|
||
|
|
existing
|
||
|
|
.as_ref()
|
||
|
|
.and_then(|existing| existing.client_family.clone())
|
||
|
|
}),
|
||
request_metadata,
|
|||
created_at_unix_ms,
|
|||
updated_at_unix_secs: usage.updated_at_unix_secs,
|
|||
|
|
finalized_at_unix_secs: usage.finalized_at_unix_secs,
|
||
|
|
};
|
||
|
|
|
||
|
|
by_request_id.insert(stored.request_id.clone(), stored.clone());
|
||
if let Some(auth_api_keys) = self.auth_api_keys.as_ref() {
|
|||
|
|
let before_contribution = existing.as_ref().and_then(api_key_usage_contribution);
|
||
|
|
let after_contribution = api_key_usage_contribution(&stored);
|
||
|
|
|
||
|
|
match (before_contribution.as_ref(), after_contribution.as_ref()) {
|
||
|
|
(Some(before), Some(after)) if before.api_key_id == after.api_key_id => {
|
||
|
|
let delta = ApiKeyUsageDelta::between(before, after);
|
||
|
|
auth_api_keys.apply_usage_stats_delta(before.api_key_id.as_str(), &delta, None);
|
||
|
|
}
|
||
|
|
_ => {
|
||
|
|
if let Some(before) = before_contribution.as_ref() {
|
||
|
|
let delta = ApiKeyUsageDelta::removal(before);
|
||
|
|
let recomputed_last_used_at_unix_secs = by_request_id
|
||
|
|
.values()
|
||
|
|
.filter_map(|item| {
|
||
|
|
item.api_key_id
|
||
|
|
.as_deref()
|
||
|
|
.filter(|api_key_id| *api_key_id == before.api_key_id.as_str())
|
||
|
|
.map(|_| item.created_at_unix_ms)
|
||
|
|
})
|
||
|
|
.max();
|
||
|
|
auth_api_keys.apply_usage_stats_delta(
|
||
|
|
before.api_key_id.as_str(),
|
||
|
|
&delta,
|
||
|
|
recomputed_last_used_at_unix_secs,
|
||
|
|
);
|
||
|
|
}
|
||
|
|
if let Some(after) = after_contribution.as_ref() {
|
||
|
|
let delta = ApiKeyUsageDelta::addition(after);
|
||
|
|
auth_api_keys.apply_usage_stats_delta(
|
||
|
|
after.api_key_id.as_str(),
|
||
|
|
&delta,
|
||
|
|
None,
|
||
|
|
);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
if let Some(provider_catalog) = self.provider_catalog.as_ref() {
|
|||
|
|
let before_contribution = existing
|
||
|
|
.as_ref()
|
||
|
|
.and_then(provider_api_key_usage_contribution);
|
||
|
|
let after_contribution = provider_api_key_usage_contribution(&stored);
|
||
|
|
|
||
|
|
match (before_contribution.as_ref(), after_contribution.as_ref()) {
|
||
|
|
(Some(before), Some(after)) if before.key_id == after.key_id => {
|
||
|
|
let delta = ProviderApiKeyUsageDelta::between(before, after);
|
||
|
|
provider_catalog.apply_usage_stats_delta(before.key_id.as_str(), &delta, None);
|
||
|
|
}
|
||
|
|
_ => {
|
||
|
|
if let Some(before) = before_contribution.as_ref() {
|
||
|
|
let delta = ProviderApiKeyUsageDelta::removal(before);
|
||
|
|
let recomputed_last_used_at_unix_secs = by_request_id
|
||
|
|
.values()
|
||
|
|
.filter_map(|item| {
|
||
|
|
item.provider_api_key_id
|
||
|
|
.as_deref()
|
||
|
|
.filter(|key_id| *key_id == before.key_id.as_str())
|
||
|
|
.map(|_| item.created_at_unix_ms)
|
||
|
|
})
|
||
|
|
.max();
|
||
|
|
provider_catalog.apply_usage_stats_delta(
|
||
|
|
before.key_id.as_str(),
|
||
|
|
&delta,
|
||
|
|
recomputed_last_used_at_unix_secs,
|
||
|
|
);
|
||
|
|
}
|
||
|
|
if let Some(after) = after_contribution.as_ref() {
|
||
|
|
let delta = ProviderApiKeyUsageDelta::addition(after);
|
||
|
|
provider_catalog.apply_usage_stats_delta(
|
||
|
|
after.key_id.as_str(),
|
||
|
|
&delta,
|
||
|
|
None,
|
||
|
|
);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
Ok(stored)
|
|||
|
|
}
|
||
|
|||
async fn rebuild_api_key_usage_stats(&self) -> Result<u64, DataLayerError> {
|
|||
|
|
let Some(auth_api_keys) = self.auth_api_keys.as_ref() else {
|
||
|
|
return Ok(0);
|
||
|
|
};
|
||
|
|
|
||
|
|
let by_request_id = self.by_request_id.read().expect("usage repository lock");
|
||
|
|
let mut aggregates = BTreeMap::new();
|
||
|
|
for usage in by_request_id.values() {
|
||
|
|
let Some(contribution) = api_key_usage_contribution(usage) else {
|
||
|
|
continue;
|
||
|
|
};
|
||
|
|
accumulate_api_key_usage_contribution(&mut aggregates, contribution);
|
||
|
|
}
|
||
|
|
auth_api_keys.rebuild_usage_stats(&aggregates);
|
||
|
|
Ok(aggregates.len() as u64)
|
||
|
|
}
|
||
|
|
|
||
async fn rebuild_provider_api_key_usage_stats(&self) -> Result<u64, DataLayerError> {
|
|||
|
|
let Some(provider_catalog) = self.provider_catalog.as_ref() else {
|
||
|
|
return Ok(0);
|
||
|
|
};
|
||
|
|
|
||
|
|
let by_request_id = self.by_request_id.read().expect("usage repository lock");
|
||
|
|
let mut aggregates = BTreeMap::new();
|
||
|
|
for usage in by_request_id.values() {
|
||
|
|
let Some(contribution) = provider_api_key_usage_contribution(usage) else {
|
||
|
|
continue;
|
||
|
|
};
|
||
|
|
accumulate_provider_api_key_usage_contribution(&mut aggregates, contribution);
|
||
|
|
}
|
||
|
|
provider_catalog.rebuild_usage_stats(&aggregates);
|
||
|
|
Ok(aggregates.len() as u64)
|
||
|
|
}
|
||
}
|
|||
|
|
|
||
|
|
#[cfg(test)]
|
||
|
|
mod tests {
|
||
use std::sync::Arc;
|
|||
|
|
|
||
use super::InMemoryUsageReadRepository;
|
|||
use crate::repository::auth::{
|
|||
|
|
AuthApiKeyReadRepository, InMemoryAuthApiKeySnapshotRepository,
|
||
|
|
StoredAuthApiKeyExportRecord, StoredAuthApiKeySnapshot,
|
||
|
|
};
|
||
use crate::repository::provider_catalog::{
|
|||
|
|
InMemoryProviderCatalogReadRepository, ProviderCatalogReadRepository,
|
||
|
|
ProviderCatalogWriteRepository, StoredProviderCatalogKey, StoredProviderCatalogProvider,
|
||
|
|
};
|
||
use crate::repository::usage::{
|
|||
|
|
StoredProviderUsageWindow, StoredRequestUsageAudit, UpsertUsageRecord, UsageReadRepository,
|
||
|
|
UsageWriteRepository,
|
||
|
|
};
|
||
use aether_data_contracts::repository::usage::{
|
|||
usage_body_ref, ProviderApiKeyWindowUsageRequest, UsageAuditAggregationGroupBy,
|
|||
UsageAuditAggregationQuery, UsageBodyField, UsageDashboardSummaryQuery,
|
|||
|
|
UsageLeaderboardGroupBy, UsageLeaderboardQuery, UsageProviderPerformanceQuery,
|
||
UsageTimeSeriesGranularity,
|
|||
};
|
|||
use serde_json::json;
|
|||
|
|||
fn sample_usage(request_id: &str, created_at_unix_ms: i64) -> StoredRequestUsageAudit {
|
|||
StoredRequestUsageAudit::new(
|
|||
|
|
"usage-1".to_string(),
|
||
|
|
request_id.to_string(),
|
||
|
|
Some("user-1".to_string()),
|
||
|
|
Some("api-key-1".to_string()),
|
||
|
|
Some("alice".to_string()),
|
||
|
|
Some("default".to_string()),
|
||
|
|
"OpenAI".to_string(),
|
||
|
|
"gpt-4.1".to_string(),
|
||
|
|
Some("gpt-4.1-mini".to_string()),
|
||
|
|
Some("provider-1".to_string()),
|
||
|
|
Some("endpoint-1".to_string()),
|
||
|
|
Some("provider-key-1".to_string()),
|
||
|
|
Some("chat".to_string()),
|
||
|
|
Some("openai:chat".to_string()),
|
||
|
|
Some("openai".to_string()),
|
||
|
|
Some("chat".to_string()),
|
||
|
|
Some("openai:chat".to_string()),
|
||
|
|
Some("openai".to_string()),
|
||
|
|
Some("chat".to_string()),
|
||
|
|
true,
|
||
|
|
false,
|
||
|
|
100,
|
||
|
|
50,
|
||
|
|
150,
|
||
|
|
0.12,
|
||
|
|
0.18,
|
||
|
|
Some(200),
|
||
|
|
None,
|
||
|
|
None,
|
||
|
|
Some(420),
|
||
|
|
Some(120),
|
||
|
|
"completed".to_string(),
|
||
|
|
"settled".to_string(),
|
||
created_at_unix_ms,
|
|||
|
|
created_at_unix_ms + 1,
|
||
|
|
Some(created_at_unix_ms + 2),
|
||
)
|
|||
|
|
.expect("usage should build")
|
||
|
|
}
|
||
|
|
|
||
fn sample_upsert_usage_record(request_id: &str) -> UpsertUsageRecord {
|
|||
|
|
UpsertUsageRecord {
|
||
|
|
request_id: request_id.to_string(),
|
||
|
|
user_id: None,
|
||
|
|
api_key_id: None,
|
||
|
|
username: None,
|
||
|
|
api_key_name: None,
|
||
|
|
provider_name: "OpenAI".to_string(),
|
||
|
|
model: "gpt-5".to_string(),
|
||
|
|
target_model: None,
|
||
|
|
provider_id: Some("provider-1".to_string()),
|
||
|
|
provider_endpoint_id: None,
|
||
|
|
provider_api_key_id: None,
|
||
|
|
request_type: None,
|
||
|
|
api_format: None,
|
||
|
|
api_family: None,
|
||
|
|
endpoint_kind: None,
|
||
|
|
endpoint_api_format: None,
|
||
|
|
provider_api_family: None,
|
||
|
|
provider_endpoint_kind: None,
|
||
|
|
has_format_conversion: Some(false),
|
||
|
|
is_stream: Some(false),
|
||
|
|
input_tokens: None,
|
||
|
|
output_tokens: None,
|
||
|
|
total_tokens: None,
|
||
|
|
cache_creation_input_tokens: None,
|
||
|
|
cache_creation_ephemeral_5m_input_tokens: None,
|
||
|
|
cache_creation_ephemeral_1h_input_tokens: None,
|
||
|
|
cache_read_input_tokens: None,
|
||
|
|
cache_creation_cost_usd: None,
|
||
|
|
cache_read_cost_usd: None,
|
||
|
|
output_price_per_1m: None,
|
||
|
|
total_cost_usd: None,
|
||
|
|
actual_total_cost_usd: None,
|
||
|
|
status_code: None,
|
||
|
|
error_message: None,
|
||
|
|
error_category: None,
|
||
|
|
response_time_ms: None,
|
||
|
|
first_byte_time_ms: None,
|
||
|
|
status: "pending".to_string(),
|
||
|
|
billing_status: "pending".to_string(),
|
||
|
|
request_headers: None,
|
||
|
|
request_body: None,
|
||
|
|
request_body_ref: None,
|
||
|
|
request_body_state: None,
|
||
|
|
provider_request_headers: None,
|
||
|
|
provider_request_body: None,
|
||
|
|
provider_request_body_ref: None,
|
||
|
|
provider_request_body_state: None,
|
||
|
|
response_headers: None,
|
||
|
|
response_body: None,
|
||
|
|
response_body_ref: None,
|
||
|
|
response_body_state: None,
|
||
|
|
client_response_headers: None,
|
||
|
|
client_response_body: None,
|
||
|
|
client_response_body_ref: None,
|
||
|
|
client_response_body_state: None,
|
||
|
|
candidate_id: None,
|
||
|
|
candidate_index: None,
|
||
|
|
key_name: None,
|
||
|
|
planner_kind: None,
|
||
|
|
route_family: None,
|
||
|
|
route_kind: None,
|
||
|
|
execution_path: None,
|
||
|
|
local_execution_runtime_miss_reason: None,
|
||
|
|
request_metadata: None,
|
||
|
|
finalized_at_unix_secs: None,
|
||
|
|
created_at_unix_ms: Some(1_700_000_000),
|
||
|
|
updated_at_unix_secs: 1_700_000_000,
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
#[tokio::test]
|
|||
|
|
async fn finds_usage_by_request_id() {
|
||
|
|
let repository = InMemoryUsageReadRepository::seed(vec![
|
||
|
|
sample_usage("req-1", 100),
|
||
|
|
sample_usage("req-2", 200),
|
||
|
|
]);
|
||
|
|
|
||
|
|
let usage = repository
|
||
|
|
.find_by_request_id("req-2")
|
||
|
|
.await
|
||
|
|
.expect("find should succeed")
|
||
|
|
.expect("usage should exist");
|
||
|
|
|
||
|
|
assert_eq!(usage.request_id, "req-2");
|
||
|
|
assert_eq!(usage.total_tokens, 150);
|
||
|
|
}
|
||
|
|||
#[tokio::test]
|
|||
|
|
async fn provider_aggregation_skips_unknown_provider_labels() {
|
||
let valid_provider = sample_usage("req-valid-provider", 300);
|
|||
|
|
|
||
|
|
let mut legacy_provider = sample_usage("req-legacy-provider", 250);
|
||
|
|
legacy_provider.provider_id = None;
|
||
|
|
legacy_provider.provider_name = "Legacy Provider".to_string();
|
||
|
|
|
||
let mut unknown = sample_usage("req-unknown-provider", 100);
|
|||
|
|
unknown.provider_id = None;
|
||
|
|
unknown.provider_name = "unknown".to_string();
|
||
|
|
|
||
|
|
let mut typo_unknown = sample_usage("req-unknow-provider", 200);
|
||
|
|
typo_unknown.provider_id = Some("unknow".to_string());
|
||
|
|
typo_unknown.provider_name = "unknow".to_string();
|
||
|
|
|
||
|
|
let repository = InMemoryUsageReadRepository::seed(vec![
|
||
valid_provider,
|
|||
|
|
legacy_provider,
|
||
unknown,
|
|||
|
|
typo_unknown,
|
||
|
|
]);
|
||
|
|
|
||
|
|
let rows = repository
|
||
|
|
.aggregate_usage_audits(&UsageAuditAggregationQuery {
|
||
|
|
created_from_unix_secs: 0,
|
||
|
|
created_until_unix_secs: 1_000,
|
||
|
|
group_by: UsageAuditAggregationGroupBy::Provider,
|
||
|
|
limit: 10,
|
||
exclude_reserved_provider_labels: false,
|
|||
})
|
|||
|
|
.await
|
||
|
|
.expect("aggregation should succeed");
|
||
|
|
|
||
assert_eq!(rows.len(), 2);
|
|||
|
|
let provider_id_row = rows
|
||
|
|
.iter()
|
||
|
|
.find(|row| row.group_key == "provider-1")
|
||
|
|
.expect("provider_id row should be present");
|
||
|
|
assert_eq!(provider_id_row.display_name.as_deref(), Some("OpenAI"));
|
||
|
|
assert_eq!(
|
||
|
|
provider_id_row.secondary_name.as_deref(),
|
||
|
|
Some("provider_id")
|
||
|
|
);
|
||
|
|
|
||
|
|
let legacy_name_row = rows
|
||
|
|
.iter()
|
||
|
|
.find(|row| row.group_key == "Legacy Provider")
|
||
|
|
.expect("legacy provider name row should be present");
|
||
|
|
assert_eq!(
|
||
|
|
legacy_name_row.display_name.as_deref(),
|
||
|
|
Some("Legacy Provider")
|
||
|
|
);
|
||
|
|
assert_eq!(
|
||
|
|
legacy_name_row.secondary_name.as_deref(),
|
||
|
|
Some("legacy_name")
|
||
|
|
);
|
||
}
|
|||
|
|
|
||
#[tokio::test]
|
|||
|
|
async fn aggregation_can_skip_unknown_provider_records_for_model_and_api_format() {
|
||
|
|
let mut unknown = sample_usage("req-unknown-provider", 100);
|
||
|
|
unknown.provider_id = None;
|
||
|
|
unknown.provider_name = "unknown".to_string();
|
||
|
|
|
||
|
|
let mut typo_unknown = sample_usage("req-unknow-provider", 200);
|
||
|
|
typo_unknown.provider_id = Some("unknow".to_string());
|
||
|
|
typo_unknown.provider_name = "unknow".to_string();
|
||
|
|
|
||
|
|
let mut pending_provider = sample_usage("req-pending-provider", 300);
|
||
|
|
pending_provider.provider_id = None;
|
||
|
|
pending_provider.provider_name = "pending".to_string();
|
||
|
|
|
||
let mut id_only_provider = sample_usage("req-id-only-provider", 350);
|
|||
|
|
id_only_provider.provider_name = "unknown".to_string();
|
||
|
|
|
||
let repository = InMemoryUsageReadRepository::seed(vec![
|
|||
|
|
sample_usage("req-valid-provider", 400),
|
||
|
|
unknown,
|
||
|
|
typo_unknown,
|
||
|
|
pending_provider,
|
||
id_only_provider,
|
|||
]);
|
|||
|
|
|
||
|
|
let model_rows = repository
|
||
|
|
.aggregate_usage_audits(&UsageAuditAggregationQuery {
|
||
|
|
created_from_unix_secs: 0,
|
||
|
|
created_until_unix_secs: 1_000,
|
||
|
|
group_by: UsageAuditAggregationGroupBy::Model,
|
||
|
|
limit: 10,
|
||
|
|
exclude_reserved_provider_labels: true,
|
||
|
|
})
|
||
|
|
.await
|
||
|
|
.expect("model aggregation should succeed");
|
||
|
|
assert_eq!(model_rows.len(), 1);
|
||
|
|
assert_eq!(model_rows[0].group_key, "gpt-4.1");
|
||
assert_eq!(model_rows[0].request_count, 2);
|
|||
|
|||
|
|
let api_format_rows = repository
|
||
|
|
.aggregate_usage_audits(&UsageAuditAggregationQuery {
|
||
|
|
created_from_unix_secs: 0,
|
||
|
|
created_until_unix_secs: 1_000,
|
||
|
|
group_by: UsageAuditAggregationGroupBy::ApiFormat,
|
||
|
|
limit: 10,
|
||
|
|
exclude_reserved_provider_labels: true,
|
||
|
|
})
|
||
|
|
.await
|
||
|
|
.expect("api format aggregation should succeed");
|
||
|
|
assert_eq!(api_format_rows.len(), 1);
|
||
|
|
assert_eq!(api_format_rows[0].group_key, "openai:chat");
|
||
assert_eq!(api_format_rows[0].request_count, 2);
|
|||
}
|
|||
|
|
|
||
#[tokio::test]
|
|||
|
|
async fn stale_pending_update_does_not_regress_finalized_usage() {
|
||
|
|
let repository = InMemoryUsageReadRepository::default();
|
||
|
|
repository
|
||
|
|
.upsert(UpsertUsageRecord {
|
||
|
|
request_id: "req-finalized-1".to_string(),
|
||
|
|
user_id: Some("user-1".to_string()),
|
||
|
|
api_key_id: Some("api-key-1".to_string()),
|
||
|
|
username: None,
|
||
|
|
api_key_name: None,
|
||
|
|
provider_name: "OpenAI".to_string(),
|
||
|
|
model: "gpt-5".to_string(),
|
||
|
|
target_model: None,
|
||
|
|
provider_id: Some("provider-1".to_string()),
|
||
|
|
provider_endpoint_id: Some("endpoint-1".to_string()),
|
||
|
|
provider_api_key_id: Some("provider-key-1".to_string()),
|
||
|
|
request_type: Some("chat".to_string()),
|
||
|
|
api_format: Some("openai:chat".to_string()),
|
||
|
|
api_family: Some("openai".to_string()),
|
||
|
|
endpoint_kind: Some("chat".to_string()),
|
||
|
|
endpoint_api_format: Some("openai:chat".to_string()),
|
||
|
|
provider_api_family: Some("openai".to_string()),
|
||
|
|
provider_endpoint_kind: Some("chat".to_string()),
|
||
|
|
has_format_conversion: Some(false),
|
||
|
|
is_stream: Some(false),
|
||
|
|
input_tokens: Some(3),
|
||
|
|
output_tokens: Some(5),
|
||
|
|
total_tokens: Some(8),
|
||
|
|
cache_creation_input_tokens: None,
|
||
|
|
cache_creation_ephemeral_5m_input_tokens: None,
|
||
|
|
cache_creation_ephemeral_1h_input_tokens: None,
|
||
|
|
cache_read_input_tokens: None,
|
||
|
|
cache_creation_cost_usd: None,
|
||
|
|
cache_read_cost_usd: None,
|
||
|
|
output_price_per_1m: None,
|
||
|
|
total_cost_usd: None,
|
||
|
|
actual_total_cost_usd: None,
|
||
|
|
status_code: Some(200),
|
||
|
|
error_message: None,
|
||
|
|
error_category: None,
|
||
|
|
response_time_ms: Some(45),
|
||
|
|
first_byte_time_ms: None,
|
||
|
|
status: "completed".to_string(),
|
||
|
|
billing_status: "pending".to_string(),
|
||
|
|
request_headers: None,
|
||
|
|
request_body: None,
|
||
|
|
request_body_ref: None,
|
||
|
|
provider_request_headers: None,
|
||
|
|
provider_request_body: None,
|
||
|
|
provider_request_body_ref: None,
|
||
|
|
response_headers: None,
|
||
|
|
response_body: None,
|
||
|
|
response_body_ref: None,
|
||
|
|
client_response_headers: None,
|
||
|
|
client_response_body: None,
|
||
|
|
client_response_body_ref: None,
|
||
request_body_state: None,
|
|||
|
|
provider_request_body_state: None,
|
||
|
|
response_body_state: None,
|
||
|
|
client_response_body_state: None,
|
||
candidate_id: None,
|
|||
|
|
candidate_index: None,
|
||
|
|
key_name: None,
|
||
|
|
planner_kind: None,
|
||
|
|
route_family: None,
|
||
|
|
route_kind: None,
|
||
|
|
execution_path: None,
|
||
|
|
local_execution_runtime_miss_reason: None,
|
||
|
|
request_metadata: None,
|
||
|
|
finalized_at_unix_secs: Some(101),
|
||
|
|
created_at_unix_ms: Some(100),
|
||
|
|
updated_at_unix_secs: 101,
|
||
|
|
})
|
||
|
|
.await
|
||
|
|
.expect("completed usage should upsert");
|
||
|
|
|
||
|
|
repository
|
||
|
|
.upsert(UpsertUsageRecord {
|
||
|
|
request_id: "req-finalized-1".to_string(),
|
||
|
|
user_id: Some("user-1".to_string()),
|
||
|
|
api_key_id: Some("api-key-1".to_string()),
|
||
|
|
username: None,
|
||
|
|
api_key_name: None,
|
||
|
|
provider_name: "OpenAI".to_string(),
|
||
|
|
model: "gpt-5".to_string(),
|
||
|
|
target_model: None,
|
||
|
|
provider_id: Some("provider-1".to_string()),
|
||
|
|
provider_endpoint_id: Some("endpoint-1".to_string()),
|
||
|
|
provider_api_key_id: Some("provider-key-1".to_string()),
|
||
|
|
request_type: Some("chat".to_string()),
|
||
|
|
api_format: Some("openai:chat".to_string()),
|
||
|
|
api_family: Some("openai".to_string()),
|
||
|
|
endpoint_kind: Some("chat".to_string()),
|
||
|
|
endpoint_api_format: Some("openai:chat".to_string()),
|
||
|
|
provider_api_family: Some("openai".to_string()),
|
||
|
|
provider_endpoint_kind: Some("chat".to_string()),
|
||
|
|
has_format_conversion: Some(false),
|
||
|
|
is_stream: Some(false),
|
||
|
|
input_tokens: None,
|
||
|
|
output_tokens: None,
|
||
|
|
total_tokens: None,
|
||
|
|
cache_creation_input_tokens: None,
|
||
|
|
cache_creation_ephemeral_5m_input_tokens: None,
|
||
|
|
cache_creation_ephemeral_1h_input_tokens: None,
|
||
|
|
cache_read_input_tokens: None,
|
||
|
|
cache_creation_cost_usd: None,
|
||
|
|
cache_read_cost_usd: None,
|
||
|
|
output_price_per_1m: None,
|
||
|
|
total_cost_usd: None,
|
||
|
|
actual_total_cost_usd: None,
|
||
|
|
status_code: None,
|
||
|
|
error_message: None,
|
||
|
|
error_category: None,
|
||
|
|
response_time_ms: None,
|
||
|
|
first_byte_time_ms: None,
|
||
|
|
status: "pending".to_string(),
|
||
|
|
billing_status: "pending".to_string(),
|
||
|
|
request_headers: None,
|
||
|
|
request_body: None,
|
||
|
|
request_body_ref: None,
|
||
|
|
provider_request_headers: None,
|
||
|
|
provider_request_body: None,
|
||
|
|
provider_request_body_ref: None,
|
||
|
|
response_headers: None,
|
||
|
|
response_body: None,
|
||
|
|
response_body_ref: None,
|
||
|
|
client_response_headers: None,
|
||
|
|
client_response_body: None,
|
||
|
|
client_response_body_ref: None,
|
||
request_body_state: None,
|
|||
|
|
provider_request_body_state: None,
|
||
|
|
response_body_state: None,
|
||
|
|
client_response_body_state: None,
|
||
candidate_id: None,
|
|||
|
|
candidate_index: None,
|
||
|
|
key_name: None,
|
||
|
|
planner_kind: None,
|
||
|
|
route_family: None,
|
||
|
|
route_kind: None,
|
||
|
|
execution_path: None,
|
||
|
|
local_execution_runtime_miss_reason: None,
|
||
|
|
request_metadata: None,
|
||
|
|
finalized_at_unix_secs: None,
|
||
|
|
created_at_unix_ms: Some(100),
|
||
|
|
updated_at_unix_secs: 102,
|
||
|
|
})
|
||
|
|
.await
|
||
|
|
.expect("stale pending usage should upsert");
|
||
|
|
|
||
|
|
let stored = repository
|
||
|
|
.find_by_request_id("req-finalized-1")
|
||
|
|
.await
|
||
|
|
.expect("usage lookup should succeed")
|
||
|
|
.expect("usage should exist");
|
||
|
|
assert_eq!(stored.status, "completed");
|
||
|
|
assert_eq!(stored.status_code, Some(200));
|
||
|
|
assert_eq!(stored.total_tokens, 8);
|
||
|
|
assert_eq!(stored.finalized_at_unix_secs, Some(101));
|
||
|
|
}
|
||
|
|
|
||
#[tokio::test]
|
|||
|
|
async fn upsert_allows_streaming_recovery_after_void_failure() {
|
||
|
|
let repository = InMemoryUsageReadRepository::default();
|
||
|
|
repository
|
||
|
|
.upsert(UpsertUsageRecord {
|
||
|
|
request_id: "req-recover-1".to_string(),
|
||
|
|
user_id: Some("user-1".to_string()),
|
||
|
|
api_key_id: Some("api-key-1".to_string()),
|
||
|
|
username: None,
|
||
|
|
api_key_name: None,
|
||
|
|
provider_name: "OpenAI".to_string(),
|
||
|
|
model: "gpt-5".to_string(),
|
||
|
|
target_model: None,
|
||
|
|
provider_id: Some("provider-1".to_string()),
|
||
|
|
provider_endpoint_id: Some("endpoint-1".to_string()),
|
||
|
|
provider_api_key_id: Some("provider-key-1".to_string()),
|
||
|
|
request_type: Some("chat".to_string()),
|
||
|
|
api_format: Some("openai:chat".to_string()),
|
||
|
|
api_family: Some("openai".to_string()),
|
||
|
|
endpoint_kind: Some("chat".to_string()),
|
||
|
|
endpoint_api_format: Some("openai:chat".to_string()),
|
||
|
|
provider_api_family: Some("openai".to_string()),
|
||
|
|
provider_endpoint_kind: Some("chat".to_string()),
|
||
|
|
has_format_conversion: Some(false),
|
||
|
|
is_stream: Some(false),
|
||
|
|
input_tokens: None,
|
||
|
|
output_tokens: None,
|
||
|
|
total_tokens: None,
|
||
|
|
cache_creation_input_tokens: None,
|
||
|
|
cache_creation_ephemeral_5m_input_tokens: None,
|
||
|
|
cache_creation_ephemeral_1h_input_tokens: None,
|
||
|
|
cache_read_input_tokens: None,
|
||
|
|
cache_creation_cost_usd: None,
|
||
|
|
cache_read_cost_usd: None,
|
||
|
|
output_price_per_1m: None,
|
||
|
|
total_cost_usd: Some(0.0),
|
||
|
|
actual_total_cost_usd: Some(0.0),
|
||
|
|
status_code: Some(503),
|
||
|
|
error_message: Some("provider timeout".to_string()),
|
||
|
|
error_category: Some("provider_error".to_string()),
|
||
|
|
response_time_ms: Some(90),
|
||
|
|
first_byte_time_ms: None,
|
||
|
|
status: "failed".to_string(),
|
||
|
|
billing_status: "void".to_string(),
|
||
|
|
request_headers: None,
|
||
|
|
request_body: None,
|
||
|
|
request_body_ref: None,
|
||
|
|
provider_request_headers: None,
|
||
|
|
provider_request_body: None,
|
||
|
|
provider_request_body_ref: None,
|
||
|
|
response_headers: None,
|
||
|
|
response_body: None,
|
||
|
|
response_body_ref: None,
|
||
|
|
client_response_headers: None,
|
||
|
|
client_response_body: None,
|
||
|
|
client_response_body_ref: None,
|
||
request_body_state: None,
|
|||
|
|
provider_request_body_state: None,
|
||
|
|
response_body_state: None,
|
||
|
|
client_response_body_state: None,
|
||
candidate_id: None,
|
|||
|
|
candidate_index: None,
|
||
|
|
key_name: None,
|
||
|
|
planner_kind: None,
|
||
|
|
route_family: None,
|
||
|
|
route_kind: None,
|
||
|
|
execution_path: None,
|
||
|
|
local_execution_runtime_miss_reason: None,
|
||
|
|
request_metadata: None,
|
||
|
|
finalized_at_unix_secs: Some(101),
|
||
|
|
created_at_unix_ms: Some(100),
|
||
|
|
updated_at_unix_secs: 101,
|
||
|
|
})
|
||
|
|
.await
|
||
|
|
.expect("failed usage should upsert");
|
||
|
|
|
||
|
|
repository
|
||
|
|
.upsert(UpsertUsageRecord {
|
||
|
|
request_id: "req-recover-1".to_string(),
|
||
|
|
user_id: Some("user-1".to_string()),
|
||
|
|
api_key_id: Some("api-key-1".to_string()),
|
||
|
|
username: None,
|
||
|
|
api_key_name: None,
|
||
|
|
provider_name: "OpenAI".to_string(),
|
||
|
|
model: "gpt-5".to_string(),
|
||
|
|
target_model: Some("gpt-5-mini".to_string()),
|
||
|
|
provider_id: Some("provider-1".to_string()),
|
||
|
|
provider_endpoint_id: Some("endpoint-1".to_string()),
|
||
|
|
provider_api_key_id: Some("provider-key-1".to_string()),
|
||
|
|
request_type: Some("chat".to_string()),
|
||
|
|
api_format: Some("openai:chat".to_string()),
|
||
|
|
api_family: Some("openai".to_string()),
|
||
|
|
endpoint_kind: Some("chat".to_string()),
|
||
|
|
endpoint_api_format: Some("openai:chat".to_string()),
|
||
|
|
provider_api_family: Some("openai".to_string()),
|
||
|
|
provider_endpoint_kind: Some("chat".to_string()),
|
||
|
|
has_format_conversion: Some(true),
|
||
|
|
is_stream: Some(true),
|
||
|
|
input_tokens: Some(10),
|
||
|
|
output_tokens: None,
|
||
|
|
total_tokens: None,
|
||
|
|
cache_creation_input_tokens: None,
|
||
|
|
cache_creation_ephemeral_5m_input_tokens: None,
|
||
|
|
cache_creation_ephemeral_1h_input_tokens: None,
|
||
|
|
cache_read_input_tokens: None,
|
||
|
|
cache_creation_cost_usd: None,
|
||
|
|
cache_read_cost_usd: None,
|
||
|
|
output_price_per_1m: None,
|
||
|
|
total_cost_usd: None,
|
||
|
|
actual_total_cost_usd: None,
|
||
|
|
status_code: None,
|
||
|
|
error_message: None,
|
||
|
|
error_category: None,
|
||
|
|
response_time_ms: Some(45),
|
||
|
|
first_byte_time_ms: Some(12),
|
||
|
|
status: "streaming".to_string(),
|
||
|
|
billing_status: "pending".to_string(),
|
||
|
|
request_headers: None,
|
||
|
|
request_body: None,
|
||
|
|
request_body_ref: None,
|
||
|
|
provider_request_headers: None,
|
||
|
|
provider_request_body: None,
|
||
|
|
provider_request_body_ref: None,
|
||
|
|
response_headers: None,
|
||
|
|
response_body: None,
|
||
|
|
response_body_ref: None,
|
||
|
|
client_response_headers: None,
|
||
|
|
client_response_body: None,
|
||
|
|
client_response_body_ref: None,
|
||
request_body_state: None,
|
|||
|
|
provider_request_body_state: None,
|
||
|
|
response_body_state: None,
|
||
|
|
client_response_body_state: None,
|
||
candidate_id: Some("cand-1".to_string()),
|
|||
|
|
candidate_index: Some(1),
|
||
|
|
key_name: Some("primary".to_string()),
|
||
|
|
planner_kind: Some("claude_cli_sync".to_string()),
|
||
|
|
route_family: Some("claude".to_string()),
|
||
|
|
route_kind: Some("cli".to_string()),
|
||
|
|
execution_path: Some("remote".to_string()),
|
||
|
|
local_execution_runtime_miss_reason: None,
|
||
|
|
request_metadata: Some(json!({
|
||
|
|
"trace_id": "trace-recovered"
|
||
|
|
})),
|
||
|
|
finalized_at_unix_secs: None,
|
||
|
|
created_at_unix_ms: Some(100),
|
||
|
|
updated_at_unix_secs: 102,
|
||
|
|
})
|
||
|
|
.await
|
||
|
|
.expect("recovery usage should upsert");
|
||
|
|
|
||
|
|
let stored = repository
|
||
|
|
.find_by_request_id("req-recover-1")
|
||
|
|
.await
|
||
|
|
.expect("usage lookup should succeed")
|
||
|
|
.expect("usage should exist");
|
||
|
|
assert_eq!(stored.status, "streaming");
|
||
|
|
assert_eq!(stored.billing_status, "pending");
|
||
|
|
assert_eq!(stored.status_code, None);
|
||
|
|
assert_eq!(stored.error_message, None);
|
||
|
|
assert_eq!(stored.finalized_at_unix_secs, None);
|
||
|
|
assert_eq!(
|
||
|
|
stored.request_metadata,
|
||
|
|
Some(json!({ "trace_id": "trace-recovered" }))
|
||
|
|
);
|
||
|
|
assert_eq!(stored.total_tokens, 10);
|
||
|
|
}
|
||
|
|
|
||
#[tokio::test]
|
|||
|
|
async fn stale_pending_update_does_not_reopen_void_failure() {
|
||
|
|
let repository = InMemoryUsageReadRepository::default();
|
||
|
|
repository
|
||
|
|
.upsert(UpsertUsageRecord {
|
||
|
|
status: "failed".to_string(),
|
||
|
|
billing_status: "void".to_string(),
|
||
|
|
status_code: Some(503),
|
||
|
|
error_message: Some("provider timeout".to_string()),
|
||
|
|
error_category: Some("provider_error".to_string()),
|
||
|
|
response_time_ms: Some(90),
|
||
|
|
finalized_at_unix_secs: Some(101),
|
||
|
|
created_at_unix_ms: Some(100),
|
||
|
|
updated_at_unix_secs: 101,
|
||
|
|
..sample_upsert_usage_record("req-void-failure-1")
|
||
|
|
})
|
||
|
|
.await
|
||
|
|
.expect("failed usage should upsert");
|
||
|
|
|
||
|
|
repository
|
||
|
|
.upsert(UpsertUsageRecord {
|
||
|
|
created_at_unix_ms: Some(100),
|
||
|
|
updated_at_unix_secs: 102,
|
||
|
|
..sample_upsert_usage_record("req-void-failure-1")
|
||
|
|
})
|
||
|
|
.await
|
||
|
|
.expect("stale pending usage should upsert");
|
||
|
|
|
||
|
|
let stored = repository
|
||
|
|
.find_by_request_id("req-void-failure-1")
|
||
|
|
.await
|
||
|
|
.expect("usage lookup should succeed")
|
||
|
|
.expect("usage should exist");
|
||
|
|
assert_eq!(stored.status, "failed");
|
||
|
|
assert_eq!(stored.billing_status, "void");
|
||
|
|
assert_eq!(stored.status_code, Some(503));
|
||
|
|
assert_eq!(stored.finalized_at_unix_secs, Some(101));
|
||
|
|
}
|
||
|
|
|
||
#[tokio::test]
|
|||
|
|
async fn stale_pending_update_does_not_regress_streaming_usage() {
|
||
|
|
let repository = InMemoryUsageReadRepository::default();
|
||
|
|
repository
|
||
|
|
.upsert(UpsertUsageRecord {
|
||
|
|
request_id: "req-streaming-1".to_string(),
|
||
|
|
user_id: Some("user-1".to_string()),
|
||
|
|
api_key_id: Some("api-key-1".to_string()),
|
||
|
|
username: None,
|
||
|
|
api_key_name: None,
|
||
|
|
provider_name: "OpenAI".to_string(),
|
||
|
|
model: "gpt-5".to_string(),
|
||
|
|
target_model: Some("gpt-5-upstream".to_string()),
|
||
|
|
provider_id: Some("provider-1".to_string()),
|
||
|
|
provider_endpoint_id: Some("endpoint-1".to_string()),
|
||
|
|
provider_api_key_id: Some("provider-key-1".to_string()),
|
||
|
|
request_type: Some("chat".to_string()),
|
||
|
|
api_format: Some("openai:chat".to_string()),
|
||
|
|
api_family: Some("openai".to_string()),
|
||
|
|
endpoint_kind: Some("chat".to_string()),
|
||
|
|
endpoint_api_format: Some("openai:chat".to_string()),
|
||
|
|
provider_api_family: Some("openai".to_string()),
|
||
|
|
provider_endpoint_kind: Some("chat".to_string()),
|
||
|
|
has_format_conversion: Some(false),
|
||
|
|
is_stream: Some(true),
|
||
|
|
input_tokens: Some(10),
|
||
|
|
output_tokens: Some(2),
|
||
|
|
total_tokens: Some(12),
|
||
|
|
cache_creation_input_tokens: None,
|
||
|
|
cache_creation_ephemeral_5m_input_tokens: None,
|
||
|
|
cache_creation_ephemeral_1h_input_tokens: None,
|
||
|
|
cache_read_input_tokens: None,
|
||
|
|
cache_creation_cost_usd: None,
|
||
|
|
cache_read_cost_usd: None,
|
||
|
|
output_price_per_1m: None,
|
||
|
|
total_cost_usd: Some(0.0),
|
||
|
|
actual_total_cost_usd: Some(0.0),
|
||
|
|
status_code: Some(200),
|
||
|
|
error_message: None,
|
||
|
|
error_category: None,
|
||
|
|
response_time_ms: Some(45),
|
||
|
|
first_byte_time_ms: Some(12),
|
||
|
|
status: "streaming".to_string(),
|
||
|
|
billing_status: "pending".to_string(),
|
||
|
|
request_headers: None,
|
||
|
|
request_body: None,
|
||
|
|
request_body_ref: None,
|
||
|
|
provider_request_headers: None,
|
||
|
|
provider_request_body: None,
|
||
|
|
provider_request_body_ref: None,
|
||
|
|
response_headers: None,
|
||
|
|
response_body: None,
|
||
|
|
response_body_ref: None,
|
||
|
|
client_response_headers: None,
|
||
|
|
client_response_body: None,
|
||
|
|
client_response_body_ref: None,
|
||
request_body_state: None,
|
|||
|
|
provider_request_body_state: None,
|
||
|
|
response_body_state: None,
|
||
|
|
client_response_body_state: None,
|
||
candidate_id: Some("cand-1".to_string()),
|
|||
|
|
candidate_index: Some(1),
|
||
|
|
key_name: Some("primary".to_string()),
|
||
|
|
planner_kind: Some("claude_cli_sync".to_string()),
|
||
|
|
route_family: Some("claude".to_string()),
|
||
|
|
route_kind: Some("cli".to_string()),
|
||
|
|
execution_path: Some("remote".to_string()),
|
||
|
|
local_execution_runtime_miss_reason: None,
|
||
|
|
request_metadata: Some(json!({
|
||
|
|
"trace_id": "trace-streaming"
|
||
|
|
})),
|
||
|
|
finalized_at_unix_secs: None,
|
||
|
|
created_at_unix_ms: Some(100),
|
||
|
|
updated_at_unix_secs: 101,
|
||
|
|
})
|
||
|
|
.await
|
||
|
|
.expect("streaming usage should upsert");
|
||
|
|
|
||
|
|
repository
|
||
|
|
.upsert(UpsertUsageRecord {
|
||
|
|
request_id: "req-streaming-1".to_string(),
|
||
|
|
user_id: Some("user-1".to_string()),
|
||
|
|
api_key_id: Some("api-key-1".to_string()),
|
||
|
|
username: None,
|
||
|
|
api_key_name: None,
|
||
|
|
provider_name: "OpenAI".to_string(),
|
||
|
|
model: "gpt-5".to_string(),
|
||
|
|
target_model: None,
|
||
|
|
provider_id: Some("provider-1".to_string()),
|
||
|
|
provider_endpoint_id: Some("endpoint-1".to_string()),
|
||
|
|
provider_api_key_id: Some("provider-key-1".to_string()),
|
||
|
|
request_type: Some("chat".to_string()),
|
||
|
|
api_format: Some("openai:chat".to_string()),
|
||
|
|
api_family: Some("openai".to_string()),
|
||
|
|
endpoint_kind: Some("chat".to_string()),
|
||
|
|
endpoint_api_format: Some("openai:chat".to_string()),
|
||
|
|
provider_api_family: Some("openai".to_string()),
|
||
|
|
provider_endpoint_kind: Some("chat".to_string()),
|
||
|
|
has_format_conversion: Some(false),
|
||
|
|
is_stream: Some(true),
|
||
|
|
input_tokens: None,
|
||
|
|
output_tokens: None,
|
||
|
|
total_tokens: None,
|
||
|
|
cache_creation_input_tokens: None,
|
||
|
|
cache_creation_ephemeral_5m_input_tokens: None,
|
||
|
|
cache_creation_ephemeral_1h_input_tokens: None,
|
||
|
|
cache_read_input_tokens: None,
|
||
|
|
cache_creation_cost_usd: None,
|
||
|
|
cache_read_cost_usd: None,
|
||
|
|
output_price_per_1m: None,
|
||
|
|
total_cost_usd: None,
|
||
|
|
actual_total_cost_usd: None,
|
||
|
|
status_code: None,
|
||
|
|
error_message: None,
|
||
|
|
error_category: None,
|
||
|
|
response_time_ms: None,
|
||
|
|
first_byte_time_ms: None,
|
||
|
|
status: "pending".to_string(),
|
||
|
|
billing_status: "pending".to_string(),
|
||
|
|
request_headers: None,
|
||
|
|
request_body: None,
|
||
|
|
request_body_ref: None,
|
||
|
|
provider_request_headers: None,
|
||
|
|
provider_request_body: None,
|
||
|
|
provider_request_body_ref: None,
|
||
|
|
response_headers: None,
|
||
|
|
response_body: None,
|
||
|
|
response_body_ref: None,
|
||
|
|
client_response_headers: None,
|
||
|
|
client_response_body: None,
|
||
|
|
client_response_body_ref: None,
|
||
request_body_state: None,
|
|||
|
|
provider_request_body_state: None,
|
||
|
|
response_body_state: None,
|
||
|
|
client_response_body_state: None,
|
||
candidate_id: None,
|
|||
|
|
candidate_index: None,
|
||
|
|
key_name: None,
|
||
|
|
planner_kind: None,
|
||
|
|
route_family: None,
|
||
|
|
route_kind: None,
|
||
|
|
execution_path: None,
|
||
|
|
local_execution_runtime_miss_reason: None,
|
||
|
|
request_metadata: None,
|
||
|
|
finalized_at_unix_secs: None,
|
||
|
|
created_at_unix_ms: Some(100),
|
||
|
|
updated_at_unix_secs: 102,
|
||
|
|
})
|
||
|
|
.await
|
||
|
|
.expect("stale pending usage should upsert");
|
||
|
|
|
||
|
|
let stored = repository
|
||
|
|
.find_by_request_id("req-streaming-1")
|
||
|
|
.await
|
||
|
|
.expect("usage lookup should succeed")
|
||
|
|
.expect("usage should exist");
|
||
|
|
assert_eq!(stored.status, "streaming");
|
||
|
|
assert_eq!(stored.status_code, Some(200));
|
||
|
|
assert_eq!(stored.first_byte_time_ms, Some(12));
|
||
|
|
assert_eq!(stored.response_time_ms, Some(45));
|
||
|
|
assert_eq!(stored.target_model.as_deref(), Some("gpt-5-upstream"));
|
||
|
|
assert_eq!(stored.total_tokens, 12);
|
||
|
|
}
|
||
|
|
|
||
#[tokio::test]
|
|||
|
|
async fn streaming_refresh_without_timing_does_not_clear_stream_timing() {
|
||
|
|
let repository = InMemoryUsageReadRepository::default();
|
||
|
|
let mut first = sample_upsert_usage_record("req-streaming-refresh");
|
||
|
|
first.is_stream = Some(true);
|
||
|
|
first.status = "streaming".to_string();
|
||
|
|
first.status_code = Some(200);
|
||
|
|
first.response_time_ms = Some(45);
|
||
|
|
first.first_byte_time_ms = Some(12);
|
||
|
|
repository
|
||
|
|
.upsert(first)
|
||
|
|
.await
|
||
|
|
.expect("streaming usage should upsert");
|
||
|
|
|
||
|
|
let mut refresh = sample_upsert_usage_record("req-streaming-refresh");
|
||
|
|
refresh.is_stream = Some(true);
|
||
|
|
refresh.status = "streaming".to_string();
|
||
refresh.status_code = None;
|
|||
repository
|
|||
|
|
.upsert(refresh)
|
||
|
|
.await
|
||
|
|
.expect("streaming refresh should upsert");
|
||
|
|
|
||
|
|
let stored = repository
|
||
|
|
.find_by_request_id("req-streaming-refresh")
|
||
|
|
.await
|
||
|
|
.expect("usage lookup should succeed")
|
||
|
|
.expect("usage should exist");
|
||
|
|
assert_eq!(stored.status, "streaming");
|
||
assert_eq!(stored.status_code, Some(200));
|
|||
assert_eq!(stored.response_time_ms, Some(45));
|
|||
|
|
assert_eq!(stored.first_byte_time_ms, Some(12));
|
||
|
|
}
|
||
|
|
|
||
#[tokio::test]
|
|||
|
|
async fn seed_hydrates_legacy_body_ref_metadata_into_typed_fields() {
|
||
|
|
let repository = InMemoryUsageReadRepository::seed(vec![StoredRequestUsageAudit {
|
||
|
|
request_metadata: Some(json!({
|
||
|
|
"request_body_ref": "usage://request/req-legacy/request_body"
|
||
|
|
})),
|
||
|
|
..sample_usage("req-legacy", 100)
|
||
|
|
}]);
|
||
|
|
|
||
|
|
let usage = repository
|
||
|
|
.find_by_request_id("req-legacy")
|
||
|
|
.await
|
||
|
|
.expect("find should succeed")
|
||
|
|
.expect("usage should exist");
|
||
|
|
|
||
|
|
assert_eq!(
|
||
|
|
usage.body_ref(UsageBodyField::RequestBody),
|
||
|
|
Some("usage://request/req-legacy/request_body")
|
||
|
|
);
|
||
|
|
assert_eq!(
|
||
|
|
usage.request_metadata,
|
||
|
|
Some(json!({
|
||
|
|
"request_body_ref": "usage://request/req-legacy/request_body"
|
||
|
|
}))
|
||
|
|
);
|
||
|
|
}
|
||
|
|
|
||
|
|
#[tokio::test]
|
||
|
|
async fn seed_ignores_invalid_or_mismatched_legacy_body_ref_metadata() {
|
||
|
|
let repository = InMemoryUsageReadRepository::seed(vec![
|
||
|
|
StoredRequestUsageAudit {
|
||
|
|
request_metadata: Some(json!({
|
||
|
|
"request_body_ref": "blob://legacy-request"
|
||
|
|
})),
|
||
|
|
..sample_usage("req-invalid-legacy", 100)
|
||
|
|
},
|
||
|
|
StoredRequestUsageAudit {
|
||
|
|
request_metadata: Some(json!({
|
||
|
|
"request_body_ref": "usage://request/req-other/request_body"
|
||
|
|
})),
|
||
|
|
..sample_usage("req-mismatched-legacy", 200)
|
||
|
|
},
|
||
|
|
]);
|
||
|
|
|
||
|
|
let invalid = repository
|
||
|
|
.find_by_request_id("req-invalid-legacy")
|
||
|
|
.await
|
||
|
|
.expect("find should succeed")
|
||
|
|
.expect("usage should exist");
|
||
|
|
let mismatched = repository
|
||
|
|
.find_by_request_id("req-mismatched-legacy")
|
||
|
|
.await
|
||
|
|
.expect("find should succeed")
|
||
|
|
.expect("usage should exist");
|
||
|
|
|
||
|
|
assert_eq!(invalid.body_ref(UsageBodyField::RequestBody), None);
|
||
|
|
assert_eq!(mismatched.body_ref(UsageBodyField::RequestBody), None);
|
||
|
|
}
|
||
|
|
|
||
|
|
#[tokio::test]
|
||
|
|
async fn detached_body_seed_moves_large_payloads_behind_usage_refs() {
|
||
|
|
let mut usage = sample_usage("req-detached", 100);
|
||
|
|
usage.request_body = Some(json!({
|
||
|
|
"model": "gpt-4.1",
|
||
|
|
"messages": [{"role": "user", "content": "hello"}]
|
||
|
|
}));
|
||
|
|
usage.provider_request_body = Some(json!({
|
||
|
|
"model": "gpt-4.1-mini",
|
||
|
|
"stream": false
|
||
|
|
}));
|
||
|
|
|
||
|
|
let repository = InMemoryUsageReadRepository::seed_with_detached_bodies(vec![usage]);
|
||
|
|
|
||
|
|
let stored = repository
|
||
|
|
.find_by_request_id("req-detached")
|
||
|
|
.await
|
||
|
|
.expect("find should succeed")
|
||
|
|
.expect("usage should exist");
|
||
|
|
|
||
|
|
assert!(stored.request_body.is_none());
|
||
|
|
assert!(stored.provider_request_body.is_none());
|
||
|
|
assert_eq!(
|
||
|
|
stored.body_ref(UsageBodyField::RequestBody),
|
||
|
|
Some("usage://request/req-detached/request_body")
|
||
|
|
);
|
||
|
|
assert_eq!(
|
||
|
|
stored.body_ref(UsageBodyField::ProviderRequestBody),
|
||
|
|
Some("usage://request/req-detached/provider_request_body")
|
||
|
|
);
|
||
|
|
assert_eq!(stored.request_metadata, None);
|
||
|
|
assert_eq!(
|
||
|
|
repository
|
||
|
|
.resolve_body_ref(&usage_body_ref("req-detached", UsageBodyField::RequestBody))
|
||
|
|
.await
|
||
|
|
.expect("body ref should resolve"),
|
||
|
|
Some(json!({
|
||
|
|
"model": "gpt-4.1",
|
||
|
|
"messages": [{"role": "user", "content": "hello"}]
|
||
|
|
}))
|
||
|
|
);
|
||
|
|
assert_eq!(
|
||
|
|
repository
|
||
|
|
.resolve_body_ref(&usage_body_ref(
|
||
|
|
"req-detached",
|
||
|
|
UsageBodyField::ProviderRequestBody
|
||
|
|
))
|
||
|
|
.await
|
||
|
|
.expect("provider request body ref should resolve"),
|
||
|
|
Some(json!({
|
||
|
|
"model": "gpt-4.1-mini",
|
||
|
|
"stream": false
|
||
|
|
}))
|
||
|
|
);
|
||
|
|
}
|
||
|
|
|
||
#[tokio::test]
|
|||
|
|
async fn upsert_writes_usage_record() {
|
||
|
|
let repository = InMemoryUsageReadRepository::default();
|
||
|
|
let stored = repository
|
||
|
|
.upsert(UpsertUsageRecord {
|
||
|
|
request_id: "req-upsert-1".to_string(),
|
||
|
|
user_id: Some("user-1".to_string()),
|
||
|
|
api_key_id: Some("key-1".to_string()),
|
||
|
|
username: None,
|
||
|
|
api_key_name: None,
|
||
|
|
provider_name: "OpenAI".to_string(),
|
||
|
|
model: "gpt-5".to_string(),
|
||
|
|
target_model: Some("gpt-5-mini".to_string()),
|
||
|
|
provider_id: Some("provider-1".to_string()),
|
||
|
|
provider_endpoint_id: Some("endpoint-1".to_string()),
|
||
|
|
provider_api_key_id: Some("provider-key-1".to_string()),
|
||
|
|
request_type: Some("chat".to_string()),
|
||
|
|
api_format: Some("openai:chat".to_string()),
|
||
|
|
api_family: Some("openai".to_string()),
|
||
|
|
endpoint_kind: Some("chat".to_string()),
|
||
|
|
endpoint_api_format: Some("openai:chat".to_string()),
|
||
|
|
provider_api_family: Some("openai".to_string()),
|
||
|
|
provider_endpoint_kind: Some("chat".to_string()),
|
||
|
|
has_format_conversion: Some(false),
|
||
|
|
is_stream: Some(true),
|
||
|
|
input_tokens: Some(10),
|
||
|
|
output_tokens: Some(20),
|
||
|
|
total_tokens: None,
|
||
|
|
cache_creation_input_tokens: None,
|
||
cache_creation_ephemeral_5m_input_tokens: None,
|
|||
|
|
cache_creation_ephemeral_1h_input_tokens: None,
|
||
cache_read_input_tokens: None,
|
|||
|
|
cache_creation_cost_usd: None,
|
||
|
|
cache_read_cost_usd: None,
|
||
|
|
output_price_per_1m: None,
|
||
|
|
total_cost_usd: Some(0.25),
|
||
|
|
actual_total_cost_usd: Some(0.15),
|
||
|
|
status_code: Some(200),
|
||
|
|
error_message: None,
|
||
|
|
error_category: None,
|
||
|
|
response_time_ms: Some(300),
|
||
|
|
first_byte_time_ms: Some(120),
|
||
|
|
status: "completed".to_string(),
|
||
|
|
billing_status: "pending".to_string(),
|
||
|
|
request_headers: Some(json!({"authorization": "Bearer test"})),
|
||
|
|
request_body: Some(json!({"model": "gpt-5"})),
|
||
request_body_ref: None,
|
|||
provider_request_headers: None,
|
|||
|
|
provider_request_body: None,
|
||
provider_request_body_ref: None,
|
|||
response_headers: None,
|
|||
|
|
response_body: None,
|
||
response_body_ref: None,
|
|||
client_response_headers: None,
|
|||
|
|
client_response_body: None,
|
||
client_response_body_ref: None,
|
|||
request_body_state: None,
|
|||
|
|
provider_request_body_state: None,
|
||
|
|
response_body_state: None,
|
||
|
|
client_response_body_state: None,
|
||
candidate_id: None,
|
|||
|
|
candidate_index: None,
|
||
|
|
key_name: None,
|
||
|
|
planner_kind: None,
|
||
|
|
route_family: None,
|
||
|
|
route_kind: None,
|
||
|
|
execution_path: None,
|
||
|
|
local_execution_runtime_miss_reason: None,
|
||
request_metadata: None,
|
|||
|
|
finalized_at_unix_secs: None,
|
||
created_at_unix_ms: Some(100),
|
|||
updated_at_unix_secs: 101,
|
|||
|
|
})
|
||
|
|
.await
|
||
|
|
.expect("upsert should succeed");
|
||
|
|
|
||
|
|
assert_eq!(stored.request_id, "req-upsert-1");
|
||
|
|
assert_eq!(stored.total_tokens, 30);
|
||
|
|
assert_eq!(stored.total_cost_usd, 0.25);
|
||
|
|
assert_eq!(stored.actual_total_cost_usd, 0.15);
|
||
|
|
assert_eq!(
|
||
|
|
repository
|
||
|
|
.find_by_request_id("req-upsert-1")
|
||
|
|
.await
|
||
|
|
.expect("find should succeed")
|
||
|
|
.expect("usage should exist")
|
||
|
|
.model,
|
||
|
|
"gpt-5"
|
||
|
|
);
|
||
|
|
}
|
||
|
|
|
||
#[tokio::test]
|
|||
|
|
async fn upsert_defaults_created_at_to_second_timestamp() {
|
||
|
|
let repository = InMemoryUsageReadRepository::default();
|
||
|
|
let stored = repository
|
||
|
|
.upsert(UpsertUsageRecord {
|
||
|
|
request_id: "req-upsert-ms-default".to_string(),
|
||
|
|
user_id: None,
|
||
|
|
api_key_id: None,
|
||
|
|
username: None,
|
||
|
|
api_key_name: None,
|
||
|
|
provider_name: "OpenAI".to_string(),
|
||
|
|
model: "gpt-5".to_string(),
|
||
|
|
target_model: None,
|
||
|
|
provider_id: None,
|
||
|
|
provider_endpoint_id: None,
|
||
|
|
provider_api_key_id: None,
|
||
|
|
request_type: None,
|
||
|
|
api_format: None,
|
||
|
|
api_family: None,
|
||
|
|
endpoint_kind: None,
|
||
|
|
endpoint_api_format: None,
|
||
|
|
provider_api_family: None,
|
||
|
|
provider_endpoint_kind: None,
|
||
|
|
has_format_conversion: None,
|
||
|
|
is_stream: None,
|
||
|
|
input_tokens: None,
|
||
|
|
output_tokens: None,
|
||
|
|
total_tokens: None,
|
||
|
|
cache_creation_input_tokens: None,
|
||
cache_creation_ephemeral_5m_input_tokens: None,
|
|||
|
|
cache_creation_ephemeral_1h_input_tokens: None,
|
||
cache_read_input_tokens: None,
|
|||
|
|
cache_creation_cost_usd: None,
|
||
|
|
cache_read_cost_usd: None,
|
||
|
|
output_price_per_1m: None,
|
||
|
|
total_cost_usd: None,
|
||
|
|
actual_total_cost_usd: None,
|
||
|
|
status_code: None,
|
||
|
|
error_message: None,
|
||
|
|
error_category: None,
|
||
|
|
response_time_ms: None,
|
||
|
|
first_byte_time_ms: None,
|
||
|
|
status: "completed".to_string(),
|
||
|
|
billing_status: "pending".to_string(),
|
||
|
|
request_headers: None,
|
||
|
|
request_body: None,
|
||
request_body_ref: None,
|
|||
provider_request_headers: None,
|
|||
|
|
provider_request_body: None,
|
||
provider_request_body_ref: None,
|
|||
response_headers: None,
|
|||
|
|
response_body: None,
|
||
response_body_ref: None,
|
|||
client_response_headers: None,
|
|||
|
|
client_response_body: None,
|
||
client_response_body_ref: None,
|
|||
request_body_state: None,
|
|||
|
|
provider_request_body_state: None,
|
||
|
|
response_body_state: None,
|
||
|
|
client_response_body_state: None,
|
||
candidate_id: None,
|
|||
|
|
candidate_index: None,
|
||
|
|
key_name: None,
|
||
|
|
planner_kind: None,
|
||
|
|
route_family: None,
|
||
|
|
route_kind: None,
|
||
|
|
execution_path: None,
|
||
|
|
local_execution_runtime_miss_reason: None,
|
||
request_metadata: None,
|
|||
|
|
finalized_at_unix_secs: None,
|
||
|
|
created_at_unix_ms: None,
|
||
|
|
updated_at_unix_secs: 101,
|
||
|
|
})
|
||
|
|
.await
|
||
|
|
.expect("upsert should succeed");
|
||
|
|
|
||
|
|
assert_eq!(stored.created_at_unix_ms, 101);
|
||
|
|
}
|
||
|
|
|
||
#[tokio::test]
|
|||
|
|
async fn upsert_does_not_backfill_legacy_output_price_from_request_metadata() {
|
||
|
|
let repository = InMemoryUsageReadRepository::default();
|
||
|
|
let stored = repository
|
||
|
|
.upsert(UpsertUsageRecord {
|
||
|
|
request_id: "req-upsert-price-metadata".to_string(),
|
||
|
|
user_id: None,
|
||
|
|
api_key_id: None,
|
||
|
|
username: None,
|
||
|
|
api_key_name: None,
|
||
|
|
provider_name: "OpenAI".to_string(),
|
||
|
|
model: "gpt-5".to_string(),
|
||
|
|
target_model: None,
|
||
|
|
provider_id: None,
|
||
|
|
provider_endpoint_id: None,
|
||
|
|
provider_api_key_id: None,
|
||
|
|
request_type: Some("chat".to_string()),
|
||
|
|
api_format: Some("openai:chat".to_string()),
|
||
|
|
api_family: Some("openai".to_string()),
|
||
|
|
endpoint_kind: Some("chat".to_string()),
|
||
|
|
endpoint_api_format: Some("openai:chat".to_string()),
|
||
|
|
provider_api_family: Some("openai".to_string()),
|
||
|
|
provider_endpoint_kind: Some("chat".to_string()),
|
||
|
|
has_format_conversion: Some(false),
|
||
|
|
is_stream: Some(false),
|
||
|
|
input_tokens: Some(10),
|
||
|
|
output_tokens: Some(20),
|
||
|
|
total_tokens: None,
|
||
|
|
cache_creation_input_tokens: None,
|
||
|
|
cache_creation_ephemeral_5m_input_tokens: None,
|
||
|
|
cache_creation_ephemeral_1h_input_tokens: None,
|
||
|
|
cache_read_input_tokens: None,
|
||
|
|
cache_creation_cost_usd: None,
|
||
|
|
cache_read_cost_usd: None,
|
||
|
|
output_price_per_1m: None,
|
||
|
|
total_cost_usd: Some(0.25),
|
||
|
|
actual_total_cost_usd: Some(0.15),
|
||
|
|
status_code: Some(200),
|
||
|
|
error_message: None,
|
||
|
|
error_category: None,
|
||
|
|
response_time_ms: Some(300),
|
||
|
|
first_byte_time_ms: Some(120),
|
||
|
|
status: "completed".to_string(),
|
||
|
|
billing_status: "pending".to_string(),
|
||
|
|
request_headers: None,
|
||
|
|
request_body: None,
|
||
|
|
request_body_ref: None,
|
||
|
|
provider_request_headers: None,
|
||
|
|
provider_request_body: None,
|
||
|
|
provider_request_body_ref: None,
|
||
|
|
response_headers: None,
|
||
|
|
response_body: None,
|
||
|
|
response_body_ref: None,
|
||
|
|
client_response_headers: None,
|
||
|
|
client_response_body: None,
|
||
|
|
client_response_body_ref: None,
|
||
request_body_state: None,
|
|||
|
|
provider_request_body_state: None,
|
||
|
|
response_body_state: None,
|
||
|
|
client_response_body_state: None,
|
||
candidate_id: None,
|
|||
|
|
candidate_index: None,
|
||
|
|
key_name: None,
|
||
|
|
planner_kind: None,
|
||
|
|
route_family: None,
|
||
|
|
route_kind: None,
|
||
|
|
execution_path: None,
|
||
|
|
local_execution_runtime_miss_reason: None,
|
||
|
|
request_metadata: Some(json!({
|
||
|
|
"output_price_per_1m": 15.0
|
||
|
|
})),
|
||
|
|
finalized_at_unix_secs: None,
|
||
|
|
created_at_unix_ms: Some(100),
|
||
|
|
updated_at_unix_secs: 101,
|
||
|
|
})
|
||
|
|
.await
|
||
|
|
.expect("upsert should succeed");
|
||
|
|
|
||
|
|
assert_eq!(stored.output_price_per_1m, None);
|
||
|
|
assert_eq!(stored.settlement_output_price_per_1m(), Some(15.0));
|
||
|
|
}
|
||
|
|
|
||
|
|
#[tokio::test]
|
||
|
|
async fn upsert_does_not_backfill_typed_body_refs_from_request_metadata() {
|
||
|
|
let repository = InMemoryUsageReadRepository::default();
|
||
|
|
let stored = repository
|
||
|
|
.upsert(UpsertUsageRecord {
|
||
|
|
request_id: "req-upsert-body-ref-metadata".to_string(),
|
||
|
|
user_id: None,
|
||
|
|
api_key_id: None,
|
||
|
|
username: None,
|
||
|
|
api_key_name: None,
|
||
|
|
provider_name: "OpenAI".to_string(),
|
||
|
|
model: "gpt-5".to_string(),
|
||
|
|
target_model: None,
|
||
|
|
provider_id: None,
|
||
|
|
provider_endpoint_id: None,
|
||
|
|
provider_api_key_id: None,
|
||
|
|
request_type: Some("chat".to_string()),
|
||
|
|
api_format: Some("openai:chat".to_string()),
|
||
|
|
api_family: Some("openai".to_string()),
|
||
|
|
endpoint_kind: Some("chat".to_string()),
|
||
|
|
endpoint_api_format: Some("openai:chat".to_string()),
|
||
|
|
provider_api_family: Some("openai".to_string()),
|
||
|
|
provider_endpoint_kind: Some("chat".to_string()),
|
||
|
|
has_format_conversion: Some(false),
|
||
|
|
is_stream: Some(false),
|
||
|
|
input_tokens: Some(10),
|
||
|
|
output_tokens: Some(20),
|
||
|
|
total_tokens: None,
|
||
|
|
cache_creation_input_tokens: None,
|
||
|
|
cache_creation_ephemeral_5m_input_tokens: None,
|
||
|
|
cache_creation_ephemeral_1h_input_tokens: None,
|
||
|
|
cache_read_input_tokens: None,
|
||
|
|
cache_creation_cost_usd: None,
|
||
|
|
cache_read_cost_usd: None,
|
||
|
|
output_price_per_1m: None,
|
||
|
|
total_cost_usd: Some(0.25),
|
||
|
|
actual_total_cost_usd: Some(0.15),
|
||
|
|
status_code: Some(200),
|
||
|
|
error_message: None,
|
||
|
|
error_category: None,
|
||
|
|
response_time_ms: Some(300),
|
||
|
|
first_byte_time_ms: Some(120),
|
||
|
|
status: "completed".to_string(),
|
||
|
|
billing_status: "pending".to_string(),
|
||
|
|
request_headers: None,
|
||
|
|
request_body: None,
|
||
|
|
request_body_ref: None,
|
||
|
|
provider_request_headers: None,
|
||
|
|
provider_request_body: None,
|
||
|
|
provider_request_body_ref: None,
|
||
|
|
response_headers: None,
|
||
|
|
response_body: None,
|
||
|
|
response_body_ref: None,
|
||
|
|
client_response_headers: None,
|
||
|
|
client_response_body: None,
|
||
|
|
client_response_body_ref: None,
|
||
request_body_state: None,
|
|||
|
|
provider_request_body_state: None,
|
||
|
|
response_body_state: None,
|
||
|
|
client_response_body_state: None,
|
||
candidate_id: None,
|
|||
|
|
candidate_index: None,
|
||
|
|
key_name: None,
|
||
|
|
planner_kind: None,
|
||
|
|
route_family: None,
|
||
|
|
route_kind: None,
|
||
|
|
execution_path: None,
|
||
|
|
local_execution_runtime_miss_reason: None,
|
||
|
|
request_metadata: Some(json!({
|
||
|
|
"request_body_ref": "usage://request/req-upsert-body-ref-metadata/request_body"
|
||
|
|
})),
|
||
|
|
finalized_at_unix_secs: None,
|
||
|
|
created_at_unix_ms: Some(100),
|
||
|
|
updated_at_unix_secs: 101,
|
||
|
|
})
|
||
|
|
.await
|
||
|
|
.expect("upsert should succeed");
|
||
|
|
|
||
|
|
assert_eq!(stored.request_body_ref, None);
|
||
|
|
assert_eq!(
|
||
|
|
stored.request_metadata,
|
||
|
|
Some(json!({
|
||
|
|
"request_body_ref": "usage://request/req-upsert-body-ref-metadata/request_body"
|
||
|
|
}))
|
||
|
|
);
|
||
|
|
}
|
||
|
|
|
||
|
|
#[tokio::test]
|
||
|
|
async fn upsert_keeps_typed_routing_fields_out_of_request_metadata() {
|
||
|
|
let repository = InMemoryUsageReadRepository::default();
|
||
|
|
let stored = repository
|
||
|
|
.upsert(UpsertUsageRecord {
|
||
|
|
request_id: "req-upsert-routing-metadata".to_string(),
|
||
|
|
user_id: None,
|
||
|
|
api_key_id: None,
|
||
|
|
username: None,
|
||
|
|
api_key_name: None,
|
||
|
|
provider_name: "OpenAI".to_string(),
|
||
|
|
model: "gpt-5".to_string(),
|
||
|
|
target_model: None,
|
||
|
|
provider_id: Some("provider-1".to_string()),
|
||
|
|
provider_endpoint_id: Some("endpoint-1".to_string()),
|
||
|
|
provider_api_key_id: Some("provider-key-1".to_string()),
|
||
|
|
request_type: Some("chat".to_string()),
|
||
|
|
api_format: Some("openai:chat".to_string()),
|
||
|
|
api_family: Some("openai".to_string()),
|
||
|
|
endpoint_kind: Some("chat".to_string()),
|
||
|
|
endpoint_api_format: Some("openai:chat".to_string()),
|
||
|
|
provider_api_family: Some("openai".to_string()),
|
||
|
|
provider_endpoint_kind: Some("chat".to_string()),
|
||
|
|
has_format_conversion: Some(true),
|
||
|
|
is_stream: Some(false),
|
||
|
|
input_tokens: Some(10),
|
||
|
|
output_tokens: Some(20),
|
||
|
|
total_tokens: None,
|
||
|
|
cache_creation_input_tokens: None,
|
||
|
|
cache_creation_ephemeral_5m_input_tokens: None,
|
||
|
|
cache_creation_ephemeral_1h_input_tokens: None,
|
||
|
|
cache_read_input_tokens: None,
|
||
|
|
cache_creation_cost_usd: None,
|
||
|
|
cache_read_cost_usd: None,
|
||
|
|
output_price_per_1m: None,
|
||
|
|
total_cost_usd: Some(0.25),
|
||
|
|
actual_total_cost_usd: Some(0.15),
|
||
|
|
status_code: Some(503),
|
||
|
|
error_message: None,
|
||
|
|
error_category: None,
|
||
|
|
response_time_ms: Some(300),
|
||
|
|
first_byte_time_ms: Some(120),
|
||
|
|
status: "failed".to_string(),
|
||
|
|
billing_status: "void".to_string(),
|
||
|
|
request_headers: None,
|
||
|
|
request_body: None,
|
||
|
|
request_body_ref: None,
|
||
|
|
provider_request_headers: None,
|
||
|
|
provider_request_body: None,
|
||
|
|
provider_request_body_ref: None,
|
||
|
|
response_headers: None,
|
||
|
|
response_body: None,
|
||
|
|
response_body_ref: None,
|
||
|
|
client_response_headers: None,
|
||
|
|
client_response_body: None,
|
||
|
|
client_response_body_ref: None,
|
||
request_body_state: None,
|
|||
|
|
provider_request_body_state: None,
|
||
|
|
response_body_state: None,
|
||
|
|
client_response_body_state: None,
|
||
candidate_id: Some("cand-typed".to_string()),
|
|||
|
|
candidate_index: Some(2),
|
||
|
|
key_name: Some("primary".to_string()),
|
||
|
|
planner_kind: Some("claude_cli_sync".to_string()),
|
||
|
|
route_family: Some("claude".to_string()),
|
||
|
|
route_kind: Some("cli".to_string()),
|
||
|
|
execution_path: Some("local_execution_runtime_miss".to_string()),
|
||
|
|
local_execution_runtime_miss_reason: Some("all_candidates_skipped".to_string()),
|
||
|
|
request_metadata: Some(json!({
|
||
|
|
"trace_id": "trace-1"
|
||
|
|
})),
|
||
|
|
finalized_at_unix_secs: None,
|
||
|
|
created_at_unix_ms: Some(100),
|
||
|
|
updated_at_unix_secs: 101,
|
||
|
|
})
|
||
|
|
.await
|
||
|
|
.expect("upsert should succeed");
|
||
|
|
|
||
|
|
assert_eq!(
|
||
|
|
stored.request_metadata,
|
||
|
|
Some(json!({ "trace_id": "trace-1" }))
|
||
|
|
);
|
||
|
|
assert_eq!(stored.routing_candidate_id(), Some("cand-typed"));
|
||
|
|
assert_eq!(stored.routing_candidate_index(), Some(2));
|
||
|
|
assert_eq!(stored.routing_key_name(), Some("primary"));
|
||
|
|
assert_eq!(stored.routing_planner_kind(), Some("claude_cli_sync"));
|
||
|
|
assert_eq!(stored.routing_route_family(), Some("claude"));
|
||
|
|
assert_eq!(stored.routing_route_kind(), Some("cli"));
|
||
|
|
assert_eq!(
|
||
|
|
stored.routing_execution_path(),
|
||
|
|
Some("local_execution_runtime_miss")
|
||
|
|
);
|
||
|
|
assert_eq!(
|
||
|
|
stored.routing_local_execution_runtime_miss_reason(),
|
||
|
|
Some("all_candidates_skipped")
|
||
|
|
);
|
||
|
|
}
|
||
|
|
|
||
|
|
#[tokio::test]
|
||
|
|
async fn upsert_does_not_persist_legacy_display_columns_for_new_rows() {
|
||
|
|
let repository = InMemoryUsageReadRepository::default();
|
||
|
|
let stored = repository
|
||
|
|
.upsert(UpsertUsageRecord {
|
||
|
|
request_id: "req-upsert-display-columns".to_string(),
|
||
|
|
user_id: Some("user-1".to_string()),
|
||
|
|
api_key_id: Some("key-1".to_string()),
|
||
|
|
username: Some("alice".to_string()),
|
||
|
|
api_key_name: Some("default".to_string()),
|
||
|
|
provider_name: "OpenAI".to_string(),
|
||
|
|
model: "gpt-5".to_string(),
|
||
|
|
target_model: None,
|
||
|
|
provider_id: None,
|
||
|
|
provider_endpoint_id: None,
|
||
|
|
provider_api_key_id: None,
|
||
|
|
request_type: Some("chat".to_string()),
|
||
|
|
api_format: Some("openai:chat".to_string()),
|
||
|
|
api_family: Some("openai".to_string()),
|
||
|
|
endpoint_kind: Some("chat".to_string()),
|
||
|
|
endpoint_api_format: Some("openai:chat".to_string()),
|
||
|
|
provider_api_family: Some("openai".to_string()),
|
||
|
|
provider_endpoint_kind: Some("chat".to_string()),
|
||
|
|
has_format_conversion: Some(false),
|
||
|
|
is_stream: Some(false),
|
||
|
|
input_tokens: Some(10),
|
||
|
|
output_tokens: Some(20),
|
||
|
|
total_tokens: None,
|
||
|
|
cache_creation_input_tokens: None,
|
||
|
|
cache_creation_ephemeral_5m_input_tokens: None,
|
||
|
|
cache_creation_ephemeral_1h_input_tokens: None,
|
||
|
|
cache_read_input_tokens: None,
|
||
|
|
cache_creation_cost_usd: None,
|
||
|
|
cache_read_cost_usd: None,
|
||
|
|
output_price_per_1m: None,
|
||
|
|
total_cost_usd: Some(0.25),
|
||
|
|
actual_total_cost_usd: Some(0.15),
|
||
|
|
status_code: Some(200),
|
||
|
|
error_message: None,
|
||
|
|
error_category: None,
|
||
|
|
response_time_ms: Some(300),
|
||
|
|
first_byte_time_ms: Some(120),
|
||
|
|
status: "completed".to_string(),
|
||
|
|
billing_status: "pending".to_string(),
|
||
|
|
request_headers: None,
|
||
|
|
request_body: None,
|
||
|
|
request_body_ref: None,
|
||
|
|
provider_request_headers: None,
|
||
|
|
provider_request_body: None,
|
||
|
|
provider_request_body_ref: None,
|
||
|
|
response_headers: None,
|
||
|
|
response_body: None,
|
||
|
|
response_body_ref: None,
|
||
|
|
client_response_headers: None,
|
||
|
|
client_response_body: None,
|
||
|
|
client_response_body_ref: None,
|
||
request_body_state: None,
|
|||
|
|
provider_request_body_state: None,
|
||
|
|
response_body_state: None,
|
||
|
|
client_response_body_state: None,
|
||
candidate_id: None,
|
|||
|
|
candidate_index: None,
|
||
|
|
key_name: None,
|
||
|
|
planner_kind: None,
|
||
|
|
route_family: None,
|
||
|
|
route_kind: None,
|
||
|
|
execution_path: None,
|
||
|
|
local_execution_runtime_miss_reason: None,
|
||
|
|
request_metadata: None,
|
||
|
|
finalized_at_unix_secs: None,
|
||
|
|
created_at_unix_ms: Some(100),
|
||
|
|
updated_at_unix_secs: 101,
|
||
|
|
})
|
||
|
|
.await
|
||
|
|
.expect("upsert should succeed");
|
||
|
|
|
||
|
|
assert_eq!(stored.username, None);
|
||
|
|
assert_eq!(stored.api_key_name, None);
|
||
|
|
}
|
||
|
|
|
||
|
|
#[tokio::test]
|
||
|
|
async fn upsert_preserves_existing_legacy_display_columns_when_new_write_omits_them() {
|
||
|
|
let repository = InMemoryUsageReadRepository::seed(vec![StoredRequestUsageAudit {
|
||
|
|
id: "usage-req-existing-display-columns".to_string(),
|
||
|
|
request_id: "req-existing-display-columns".to_string(),
|
||
|
|
user_id: Some("user-1".to_string()),
|
||
|
|
api_key_id: Some("key-1".to_string()),
|
||
|
|
username: Some("legacy-alice".to_string()),
|
||
|
|
api_key_name: Some("legacy-default".to_string()),
|
||
|
|
provider_name: "OpenAI".to_string(),
|
||
|
|
model: "gpt-5".to_string(),
|
||
|
|
target_model: None,
|
||
|
|
provider_id: None,
|
||
|
|
provider_endpoint_id: None,
|
||
|
|
provider_api_key_id: None,
|
||
|
|
request_type: Some("chat".to_string()),
|
||
|
|
api_format: Some("openai:chat".to_string()),
|
||
|
|
api_family: Some("openai".to_string()),
|
||
|
|
endpoint_kind: Some("chat".to_string()),
|
||
|
|
endpoint_api_format: Some("openai:chat".to_string()),
|
||
|
|
provider_api_family: Some("openai".to_string()),
|
||
|
|
provider_endpoint_kind: Some("chat".to_string()),
|
||
|
|
has_format_conversion: false,
|
||
|
|
is_stream: false,
|
||
client_family: None,
|
|||
input_tokens: 10,
|
|||
|
|
output_tokens: 20,
|
||
|
|
total_tokens: 30,
|
||
|
|
cache_creation_input_tokens: 0,
|
||
|
|
cache_creation_ephemeral_5m_input_tokens: 0,
|
||
|
|
cache_creation_ephemeral_1h_input_tokens: 0,
|
||
|
|
cache_read_input_tokens: 0,
|
||
|
|
cache_creation_cost_usd: 0.0,
|
||
|
|
cache_read_cost_usd: 0.0,
|
||
|
|
output_price_per_1m: None,
|
||
|
|
total_cost_usd: 0.25,
|
||
|
|
actual_total_cost_usd: 0.15,
|
||
|
|
status_code: Some(200),
|
||
|
|
error_message: None,
|
||
|
|
error_category: None,
|
||
|
|
response_time_ms: Some(300),
|
||
|
|
first_byte_time_ms: Some(120),
|
||
|
|
status: "completed".to_string(),
|
||
|
|
billing_status: "pending".to_string(),
|
||
|
|
request_headers: None,
|
||
|
|
request_body: None,
|
||
|
|
request_body_ref: None,
|
||
|
|
provider_request_headers: None,
|
||
|
|
provider_request_body: None,
|
||
|
|
provider_request_body_ref: None,
|
||
|
|
response_headers: None,
|
||
|
|
response_body: None,
|
||
|
|
response_body_ref: None,
|
||
|
|
client_response_headers: None,
|
||
|
|
client_response_body: None,
|
||
|
|
client_response_body_ref: None,
|
||
request_body_state: None,
|
|||
|
|
provider_request_body_state: None,
|
||
|
|
response_body_state: None,
|
||
|
|
client_response_body_state: None,
|
||
candidate_id: None,
|
|||
|
|
candidate_index: None,
|
||
|
|
key_name: None,
|
||
|
|
planner_kind: None,
|
||
|
|
route_family: None,
|
||
|
|
route_kind: None,
|
||
|
|
execution_path: None,
|
||
|
|
local_execution_runtime_miss_reason: None,
|
||
|
|
request_metadata: None,
|
||
|
|
created_at_unix_ms: 100,
|
||
|
|
updated_at_unix_secs: 101,
|
||
|
|
finalized_at_unix_secs: None,
|
||
|
|
}]);
|
||
|
|
let stored = repository
|
||
|
|
.upsert(UpsertUsageRecord {
|
||
|
|
request_id: "req-existing-display-columns".to_string(),
|
||
|
|
user_id: Some("user-1".to_string()),
|
||
|
|
api_key_id: Some("key-1".to_string()),
|
||
|
|
username: Some("fresh-alice".to_string()),
|
||
|
|
api_key_name: Some("fresh-default".to_string()),
|
||
|
|
provider_name: "OpenAI".to_string(),
|
||
|
|
model: "gpt-5-mini".to_string(),
|
||
|
|
target_model: None,
|
||
|
|
provider_id: None,
|
||
|
|
provider_endpoint_id: None,
|
||
|
|
provider_api_key_id: None,
|
||
|
|
request_type: Some("chat".to_string()),
|
||
|
|
api_format: Some("openai:chat".to_string()),
|
||
|
|
api_family: Some("openai".to_string()),
|
||
|
|
endpoint_kind: Some("chat".to_string()),
|
||
|
|
endpoint_api_format: Some("openai:chat".to_string()),
|
||
|
|
provider_api_family: Some("openai".to_string()),
|
||
|
|
provider_endpoint_kind: Some("chat".to_string()),
|
||
|
|
has_format_conversion: Some(false),
|
||
|
|
is_stream: Some(false),
|
||
|
|
input_tokens: Some(30),
|
||
|
|
output_tokens: Some(40),
|
||
|
|
total_tokens: None,
|
||
|
|
cache_creation_input_tokens: None,
|
||
|
|
cache_creation_ephemeral_5m_input_tokens: None,
|
||
|
|
cache_creation_ephemeral_1h_input_tokens: None,
|
||
|
|
cache_read_input_tokens: None,
|
||
|
|
cache_creation_cost_usd: None,
|
||
|
|
cache_read_cost_usd: None,
|
||
|
|
output_price_per_1m: None,
|
||
|
|
total_cost_usd: Some(0.45),
|
||
|
|
actual_total_cost_usd: Some(0.30),
|
||
|
|
status_code: Some(200),
|
||
|
|
error_message: None,
|
||
|
|
error_category: None,
|
||
|
|
response_time_ms: Some(200),
|
||
|
|
first_byte_time_ms: Some(80),
|
||
|
|
status: "completed".to_string(),
|
||
|
|
billing_status: "pending".to_string(),
|
||
|
|
request_headers: None,
|
||
|
|
request_body: None,
|
||
|
|
request_body_ref: None,
|
||
|
|
provider_request_headers: None,
|
||
|
|
provider_request_body: None,
|
||
|
|
provider_request_body_ref: None,
|
||
|
|
response_headers: None,
|
||
|
|
response_body: None,
|
||
|
|
response_body_ref: None,
|
||
|
|
client_response_headers: None,
|
||
|
|
client_response_body: None,
|
||
|
|
client_response_body_ref: None,
|
||
request_body_state: None,
|
|||
|
|
provider_request_body_state: None,
|
||
|
|
response_body_state: None,
|
||
|
|
client_response_body_state: None,
|
||
candidate_id: None,
|
|||
|
|
candidate_index: None,
|
||
|
|
key_name: None,
|
||
|
|
planner_kind: None,
|
||
|
|
route_family: None,
|
||
|
|
route_kind: None,
|
||
|
|
execution_path: None,
|
||
|
|
local_execution_runtime_miss_reason: None,
|
||
|
|
request_metadata: None,
|
||
|
|
finalized_at_unix_secs: None,
|
||
|
|
created_at_unix_ms: Some(100),
|
||
|
|
updated_at_unix_secs: 102,
|
||
|
|
})
|
||
|
|
.await
|
||
|
|
.expect("upsert should succeed");
|
||
|
|
|
||
|
|
assert_eq!(stored.username.as_deref(), Some("legacy-alice"));
|
||
|
|
assert_eq!(stored.api_key_name.as_deref(), Some("legacy-default"));
|
||
|
|
assert_eq!(stored.model, "gpt-5-mini");
|
||
|
|
}
|
||
|
|
|
||
#[tokio::test]
|
|||
|
|
async fn summarizes_provider_usage_windows_since_timestamp() {
|
||
|
|
let repository = InMemoryUsageReadRepository::default().with_provider_usage_windows(vec![
|
||
|
|
StoredProviderUsageWindow::new(
|
||
|
|
"provider-1".to_string(),
|
||
|
|
1_700_000_000,
|
||
|
|
10,
|
||
|
|
9,
|
||
|
|
1,
|
||
|
|
120.0,
|
||
|
|
1.25,
|
||
|
|
)
|
||
|
|
.expect("window should build"),
|
||
|
|
StoredProviderUsageWindow::new(
|
||
|
|
"provider-1".to_string(),
|
||
|
|
1_700_003_600,
|
||
|
|
6,
|
||
|
|
5,
|
||
|
|
1,
|
||
|
|
180.0,
|
||
|
|
0.75,
|
||
|
|
)
|
||
|
|
.expect("window should build"),
|
||
|
|
StoredProviderUsageWindow::new(
|
||
|
|
"provider-2".to_string(),
|
||
|
|
1_700_003_600,
|
||
|
|
99,
|
||
|
|
99,
|
||
|
|
0,
|
||
|
|
50.0,
|
||
|
|
5.0,
|
||
|
|
)
|
||
|
|
.expect("window should build"),
|
||
|
|
]);
|
||
|
|
|
||
|
|
let summary = repository
|
||
|
|
.summarize_provider_usage_since("provider-1", 1_700_000_100)
|
||
|
|
.await
|
||
|
|
.expect("summary should succeed");
|
||
|
|
|
||
|
|
assert_eq!(summary.total_requests, 6);
|
||
|
|
assert_eq!(summary.successful_requests, 5);
|
||
|
|
assert_eq!(summary.failed_requests, 1);
|
||
|
|
assert_eq!(summary.avg_response_time_ms, 180.0);
|
||
|
|
assert_eq!(summary.total_cost_usd, 0.75);
|
||
|
|
}
|
||
|
|||
|
|
#[tokio::test]
|
||
|
|
async fn summarizes_usage_by_provider_api_key_ids() {
|
||
|
|
let repository = InMemoryUsageReadRepository::seed(vec![
|
||
|
|
sample_usage("req-1", 1_711_000_000),
|
||
|
|
sample_usage("req-2", 1_711_000_250),
|
||
|
|
]);
|
||
|
|
|
||
|
|
let usage = repository
|
||
|
|
.summarize_usage_by_provider_api_key_ids(&["provider-key-1".to_string()])
|
||
|
|
.await
|
||
|
|
.expect("summary should succeed");
|
||
|
|
let item = usage
|
||
|
|
.get("provider-key-1")
|
||
|
|
.expect("provider key summary should exist");
|
||
|
|
|
||
|
|
assert_eq!(item.request_count, 2);
|
||
|
|
assert_eq!(item.total_tokens, 300);
|
||
|
|
assert_eq!(item.total_cost_usd, 0.24);
|
||
|
|
assert_eq!(item.last_used_at_unix_secs, Some(1_711_000_250));
|
||
|
|
}
|
||
|
|||
#[tokio::test]
|
|||
|
|
async fn summarizes_provider_api_key_window_usage_with_zero_rows() {
|
||
|
|
let repository = InMemoryUsageReadRepository::seed(vec![
|
||
|
|
sample_usage("req-1", 1_711_000_000),
|
||
|
|
sample_usage("req-2", 1_711_000_250),
|
||
|
|
]);
|
||
|
|
|
||
|
|
let usage = repository
|
||
|
|
.summarize_usage_by_provider_api_key_windows(&[
|
||
|
|
ProviderApiKeyWindowUsageRequest {
|
||
|
|
provider_api_key_id: "provider-key-1".to_string(),
|
||
|
|
window_code: "5h".to_string(),
|
||
|
|
start_unix_secs: 1_711_000_000,
|
||
|
|
end_unix_secs: 1_711_000_300,
|
||
|
|
},
|
||
|
|
ProviderApiKeyWindowUsageRequest {
|
||
|
|
provider_api_key_id: "provider-key-empty".to_string(),
|
||
|
|
window_code: "weekly".to_string(),
|
||
|
|
start_unix_secs: 1_711_000_000,
|
||
|
|
end_unix_secs: 1_711_000_300,
|
||
|
|
},
|
||
|
|
])
|
||
|
|
.await
|
||
|
|
.expect("window summary should succeed");
|
||
|
|
|
||
|
|
assert_eq!(usage.len(), 2);
|
||
|
|
assert_eq!(usage[0].provider_api_key_id, "provider-key-1");
|
||
|
|
assert_eq!(usage[0].window_code, "5h");
|
||
|
|
assert_eq!(usage[0].request_count, 2);
|
||
|
|
assert_eq!(usage[0].total_tokens, 300);
|
||
|
|
assert_eq!(usage[0].total_cost_usd, 0.24);
|
||
|
|
assert_eq!(usage[1].provider_api_key_id, "provider-key-empty");
|
||
|
|
assert_eq!(usage[1].window_code, "weekly");
|
||
|
|
assert_eq!(usage[1].request_count, 0);
|
||
|
|
assert_eq!(usage[1].total_tokens, 0);
|
||
|
|
assert_eq!(usage[1].total_cost_usd, 0.0);
|
||
|
|
}
|
||
|
|
|
||
#[tokio::test]
|
|||
async fn list_usage_audits_applies_second_based_time_filters() {
|
|||
let repository = InMemoryUsageReadRepository::seed(vec![
|
|||
sample_usage("req-1", 1),
|
|||
|
|
sample_usage("req-2", 2),
|
||
|
|
sample_usage("req-3", 3),
|
||
]);
|
|||
|
|
|
||
|
|
let items = repository
|
||
|
|
.list_usage_audits(&crate::repository::usage::UsageAuditListQuery {
|
||
|
|
created_from_unix_secs: Some(2),
|
||
|
|
created_until_unix_secs: Some(3),
|
||
|
|
..Default::default()
|
||
|
|
})
|
||
|
|
.await
|
||
|
|
.expect("list should succeed");
|
||
|
|
|
||
|
|
assert_eq!(items.len(), 1);
|
||
|
|
assert_eq!(items[0].request_id, "req-2");
|
||
|
|
}
|
||
|
|
|
||
#[tokio::test]
|
|||
|
|
async fn dashboard_and_leaderboard_total_tokens_use_effective_cache_aware_tokens() {
|
||
|
|
let mut item = sample_usage("req-cache-aware-total", 1_711_000_000);
|
||
|
|
item.input_tokens = 100;
|
||
|
|
item.output_tokens = 20;
|
||
|
|
item.total_tokens = 999;
|
||
|
|
item.cache_creation_input_tokens = 0;
|
||
|
|
item.cache_creation_ephemeral_5m_input_tokens = 12;
|
||
|
|
item.cache_creation_ephemeral_1h_input_tokens = 8;
|
||
|
|
item.cache_read_input_tokens = 80;
|
||
|
|
|
||
|
|
let repository = InMemoryUsageReadRepository::seed(vec![item]);
|
||
|
|
|
||
|
|
let dashboard = repository
|
||
|
|
.summarize_dashboard_usage(&UsageDashboardSummaryQuery {
|
||
|
|
created_from_unix_secs: 1_711_000_000,
|
||
|
|
created_until_unix_secs: 1_711_000_001,
|
||
|
|
user_id: None,
|
||
|
|
})
|
||
|
|
.await
|
||
|
|
.expect("dashboard should summarize");
|
||
assert_eq!(dashboard.effective_input_tokens, 0);
|
|||
assert_eq!(dashboard.cache_creation_tokens, 20);
|
|||
assert_eq!(dashboard.total_tokens, 120);
|
|||
|
|||
|
|
let leaderboard = repository
|
||
|
|
.summarize_usage_leaderboard(&UsageLeaderboardQuery {
|
||
|
|
created_from_unix_secs: 1_711_000_000,
|
||
|
|
created_until_unix_secs: 1_711_000_001,
|
||
|
|
group_by: UsageLeaderboardGroupBy::User,
|
||
|
|
user_id: None,
|
||
|
|
provider_name: None,
|
||
|
|
model: None,
|
||
|
|
})
|
||
|
|
.await
|
||
|
|
.expect("leaderboard should summarize");
|
||
|
|
assert_eq!(leaderboard.len(), 1);
|
||
assert_eq!(leaderboard[0].total_tokens, 120);
|
|||
}
|
|||
|
|
|
||
#[tokio::test]
|
|||
|
|
async fn summarizes_provider_api_key_last_used_at_in_seconds() {
|
||
|
|
let repository = InMemoryUsageReadRepository::seed(vec![
|
||
|
|
sample_usage("req-1", 1_999),
|
||
|
|
sample_usage("req-2", 2_500),
|
||
|
|
]);
|
||
|
|
|
||
|
|
let summary = repository
|
||
|
|
.summarize_usage_by_provider_api_key_ids(&["provider-key-1".to_string()])
|
||
|
|
.await
|
||
|
|
.expect("summary should succeed");
|
||
|
|
|
||
|
|
let usage = summary
|
||
|
|
.get("provider-key-1")
|
||
|
|
.expect("provider key summary should exist");
|
||
|
|
assert_eq!(usage.request_count, 2);
|
||
assert_eq!(usage.last_used_at_unix_secs, Some(2_500));
|
|||
}
|
|||
|
|||
|
|
fn sample_provider_catalog_key(key_id: &str) -> StoredProviderCatalogKey {
|
||
|
|
StoredProviderCatalogKey::new(
|
||
|
|
key_id.to_string(),
|
||
|
|
"provider-1".to_string(),
|
||
|
|
"provider key".to_string(),
|
||
|
|
"api_key".to_string(),
|
||
|
|
None,
|
||
|
|
true,
|
||
|
|
)
|
||
|
|
.expect("provider key should build")
|
||
|
|
}
|
||
|
|
|
||
|
|
fn sample_provider_catalog_repository(
|
||
|
|
key_ids: &[&str],
|
||
|
|
) -> Arc<InMemoryProviderCatalogReadRepository> {
|
||
|
|
Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||
|
|
vec![StoredProviderCatalogProvider::new(
|
||
|
|
"provider-1".to_string(),
|
||
|
|
"OpenAI".to_string(),
|
||
|
|
None,
|
||
|
|
"openai".to_string(),
|
||
|
|
)
|
||
|
|
.expect("provider should build")],
|
||
|
|
Vec::new(),
|
||
|
|
key_ids
|
||
|
|
.iter()
|
||
|
|
.map(|key_id| sample_provider_catalog_key(key_id))
|
||
|
|
.collect(),
|
||
|
|
))
|
||
|
|
}
|
||
|
|
|
||
fn sample_auth_api_key_repository(
|
|||
|
|
api_key_ids: &[&str],
|
||
|
|
) -> Arc<InMemoryAuthApiKeySnapshotRepository> {
|
||
|
|
let snapshots = api_key_ids.iter().map(|api_key_id| {
|
||
|
|
(
|
||
|
|
Some(format!("hash-{api_key_id}")),
|
||
|
|
StoredAuthApiKeySnapshot::new(
|
||
|
|
"user-1".to_string(),
|
||
|
|
"alice".to_string(),
|
||
|
|
Some("[email protected]".to_string()),
|
||
|
|
"user".to_string(),
|
||
|
|
"local".to_string(),
|
||
|
|
true,
|
||
|
|
false,
|
||
|
|
None,
|
||
|
|
None,
|
||
|
|
None,
|
||
|
|
(*api_key_id).to_string(),
|
||
|
|
Some(format!("Key {api_key_id}")),
|
||
|
|
true,
|
||
|
|
false,
|
||
|
|
false,
|
||
|
|
Some(120),
|
||
|
|
Some(8),
|
||
|
|
None,
|
||
|
|
None,
|
||
|
|
None,
|
||
|
|
None,
|
||
|
|
)
|
||
|
|
.expect("snapshot should build"),
|
||
|
|
)
|
||
|
|
});
|
||
|
|
let export_records = api_key_ids.iter().map(|api_key_id| {
|
||
|
|
StoredAuthApiKeyExportRecord::new(
|
||
|
|
"user-1".to_string(),
|
||
|
|
(*api_key_id).to_string(),
|
||
|
|
format!("hash-{api_key_id}"),
|
||
|
|
Some(format!("enc-{api_key_id}")),
|
||
|
|
Some(format!("Key {api_key_id}")),
|
||
|
|
None,
|
||
|
|
None,
|
||
|
|
None,
|
||
|
|
Some(120),
|
||
|
|
Some(8),
|
||
|
|
None,
|
||
|
|
true,
|
||
|
|
None,
|
||
|
|
false,
|
||
|
|
0,
|
||
|
|
0,
|
||
|
|
0.0,
|
||
|
|
false,
|
||
|
|
)
|
||
|
|
.expect("export record should build")
|
||
|
|
});
|
||
|
|
Arc::new(
|
||
|
|
InMemoryAuthApiKeySnapshotRepository::seed(snapshots)
|
||
|
|
.with_export_records(export_records),
|
||
|
|
)
|
||
|
|
}
|
||
|
|
|
||
#[tokio::test]
|
|||
|
|
async fn upsert_syncs_linked_provider_key_stats_without_double_counting_request_count() {
|
||
|
|
let provider_catalog = sample_provider_catalog_repository(&["provider-key-1"]);
|
||
|
|
let repository = InMemoryUsageReadRepository::default()
|
||
|
|
.with_provider_catalog_repository(Arc::clone(&provider_catalog));
|
||
|
|
|
||
|
|
repository
|
||
|
|
.upsert(UpsertUsageRecord {
|
||
|
|
provider_api_key_id: Some("provider-key-1".to_string()),
|
||
|
|
total_tokens: Some(100),
|
||
|
|
total_cost_usd: Some(0.5),
|
||
|
|
created_at_unix_ms: Some(1_711_100_000),
|
||
|
|
updated_at_unix_secs: 1_711_100_000,
|
||
|
|
..sample_upsert_usage_record("req-linked-1")
|
||
|
|
})
|
||
|
|
.await
|
||
|
|
.expect("pending upsert should succeed");
|
||
|
|
repository
|
||
|
|
.upsert(UpsertUsageRecord {
|
||
|
|
provider_api_key_id: Some("provider-key-1".to_string()),
|
||
|
|
status: "completed".to_string(),
|
||
|
|
billing_status: "settled".to_string(),
|
||
|
|
status_code: Some(200),
|
||
|
|
response_time_ms: Some(240),
|
||
|
|
total_tokens: Some(180),
|
||
|
|
total_cost_usd: Some(0.75),
|
||
|
|
created_at_unix_ms: Some(1_711_100_000),
|
||
|
|
updated_at_unix_secs: 1_711_100_010,
|
||
|
|
finalized_at_unix_secs: Some(1_711_100_011),
|
||
|
|
..sample_upsert_usage_record("req-linked-1")
|
||
|
|
})
|
||
|
|
.await
|
||
|
|
.expect("completed upsert should succeed");
|
||
|
|
|
||
|
|
let key = provider_catalog
|
||
|
|
.list_keys_by_ids(&["provider-key-1".to_string()])
|
||
|
|
.await
|
||
|
|
.expect("key list should succeed")
|
||
|
|
.into_iter()
|
||
|
|
.next()
|
||
|
|
.expect("provider key should exist");
|
||
|
|
assert_eq!(key.request_count, Some(1));
|
||
|
|
assert_eq!(key.success_count, Some(1));
|
||
|
|
assert_eq!(key.error_count, Some(0));
|
||
|
|
assert_eq!(key.total_tokens, 180);
|
||
|
|
assert_eq!(key.total_cost_usd, 0.75);
|
||
|
|
assert_eq!(key.total_response_time_ms, Some(240));
|
||
|
|
assert_eq!(key.last_used_at_unix_secs, Some(1_711_100_000));
|
||
|
|
}
|
||
|
|
|
||
#[tokio::test]
|
|||
|
|
async fn upsert_syncs_linked_api_key_stats_without_double_counting_request_count() {
|
||
|
|
let auth_api_keys = sample_auth_api_key_repository(&["api-key-1"]);
|
||
|
|
let repository = InMemoryUsageReadRepository::default()
|
||
|
|
.with_auth_api_key_repository(Arc::clone(&auth_api_keys));
|
||
|
|
|
||
|
|
repository
|
||
|
|
.upsert(UpsertUsageRecord {
|
||
|
|
api_key_id: Some("api-key-1".to_string()),
|
||
|
|
total_tokens: Some(100),
|
||
|
|
total_cost_usd: Some(0.5),
|
||
|
|
created_at_unix_ms: Some(1_711_100_000),
|
||
|
|
updated_at_unix_secs: 1_711_100_000,
|
||
|
|
..sample_upsert_usage_record("req-api-key-1")
|
||
|
|
})
|
||
|
|
.await
|
||
|
|
.expect("pending upsert should succeed");
|
||
|
|
repository
|
||
|
|
.upsert(UpsertUsageRecord {
|
||
|
|
api_key_id: Some("api-key-1".to_string()),
|
||
|
|
status: "completed".to_string(),
|
||
|
|
billing_status: "settled".to_string(),
|
||
|
|
total_tokens: Some(180),
|
||
|
|
total_cost_usd: Some(0.75),
|
||
|
|
created_at_unix_ms: Some(1_711_100_000),
|
||
|
|
updated_at_unix_secs: 1_711_100_010,
|
||
|
|
finalized_at_unix_secs: Some(1_711_100_011),
|
||
|
|
..sample_upsert_usage_record("req-api-key-1")
|
||
|
|
})
|
||
|
|
.await
|
||
|
|
.expect("completed upsert should succeed");
|
||
|
|
|
||
|
|
let key = auth_api_keys
|
||
|
|
.list_export_api_keys_by_ids(&["api-key-1".to_string()])
|
||
|
|
.await
|
||
|
|
.expect("key list should succeed")
|
||
|
|
.into_iter()
|
||
|
|
.next()
|
||
|
|
.expect("api key should exist");
|
||
|
|
assert_eq!(key.total_requests, 1);
|
||
|
|
assert_eq!(key.total_tokens, 180);
|
||
|
|
assert_eq!(key.total_cost_usd, 0.75);
|
||
|
|
}
|
||
|
|
|
||
#[tokio::test]
|
|||
|
|
async fn upsert_moves_linked_provider_key_stats_when_key_assignment_changes() {
|
||
|
|
let provider_catalog =
|
||
|
|
sample_provider_catalog_repository(&["provider-key-a", "provider-key-b"]);
|
||
|
|
let repository = InMemoryUsageReadRepository::default()
|
||
|
|
.with_provider_catalog_repository(Arc::clone(&provider_catalog));
|
||
|
|
|
||
|
|
repository
|
||
|
|
.upsert(UpsertUsageRecord {
|
||
|
|
provider_api_key_id: Some("provider-key-a".to_string()),
|
||
|
|
status: "completed".to_string(),
|
||
|
|
billing_status: "settled".to_string(),
|
||
|
|
status_code: Some(200),
|
||
|
|
response_time_ms: Some(100),
|
||
|
|
total_tokens: Some(120),
|
||
|
|
total_cost_usd: Some(0.4),
|
||
|
|
created_at_unix_ms: Some(1_711_200_000),
|
||
|
|
updated_at_unix_secs: 1_711_200_000,
|
||
|
|
finalized_at_unix_secs: Some(1_711_200_001),
|
||
|
|
..sample_upsert_usage_record("req-move-1")
|
||
|
|
})
|
||
|
|
.await
|
||
|
|
.expect("first upsert should succeed");
|
||
|
|
repository
|
||
|
|
.upsert(UpsertUsageRecord {
|
||
|
|
provider_api_key_id: Some("provider-key-b".to_string()),
|
||
|
|
status: "completed".to_string(),
|
||
|
|
billing_status: "settled".to_string(),
|
||
|
|
status_code: Some(200),
|
||
|
|
response_time_ms: Some(150),
|
||
|
|
total_tokens: Some(140),
|
||
|
|
total_cost_usd: Some(0.6),
|
||
|
|
created_at_unix_ms: Some(1_711_200_000),
|
||
|
|
updated_at_unix_secs: 1_711_200_010,
|
||
|
|
finalized_at_unix_secs: Some(1_711_200_011),
|
||
|
|
..sample_upsert_usage_record("req-move-1")
|
||
|
|
})
|
||
|
|
.await
|
||
|
|
.expect("moved upsert should succeed");
|
||
|
|
|
||
|
|
let keys = provider_catalog
|
||
|
|
.list_keys_by_ids(&["provider-key-a".to_string(), "provider-key-b".to_string()])
|
||
|
|
.await
|
||
|
|
.expect("key list should succeed");
|
||
|
|
let key_a = keys
|
||
|
|
.iter()
|
||
|
|
.find(|key| key.id == "provider-key-a")
|
||
|
|
.expect("key a should exist");
|
||
|
|
let key_b = keys
|
||
|
|
.iter()
|
||
|
|
.find(|key| key.id == "provider-key-b")
|
||
|
|
.expect("key b should exist");
|
||
|
|
|
||
|
|
assert_eq!(key_a.request_count, Some(0));
|
||
|
|
assert_eq!(key_a.success_count, Some(0));
|
||
|
|
assert_eq!(key_a.total_tokens, 0);
|
||
|
|
assert_eq!(key_a.total_cost_usd, 0.0);
|
||
|
|
assert_eq!(key_a.total_response_time_ms, Some(0));
|
||
|
|
assert_eq!(key_a.last_used_at_unix_secs, None);
|
||
|
|
|
||
|
|
assert_eq!(key_b.request_count, Some(1));
|
||
|
|
assert_eq!(key_b.success_count, Some(1));
|
||
|
|
assert_eq!(key_b.total_tokens, 140);
|
||
|
|
assert_eq!(key_b.total_cost_usd, 0.6);
|
||
|
|
assert_eq!(key_b.total_response_time_ms, Some(150));
|
||
|
|
assert_eq!(key_b.last_used_at_unix_secs, Some(1_711_200_000));
|
||
|
|
}
|
||
|
|
|
||
#[tokio::test]
|
|||
|
|
async fn upsert_moves_linked_api_key_stats_when_key_assignment_changes() {
|
||
|
|
let auth_api_keys = sample_auth_api_key_repository(&["api-key-a", "api-key-b"]);
|
||
|
|
let repository = InMemoryUsageReadRepository::default()
|
||
|
|
.with_auth_api_key_repository(Arc::clone(&auth_api_keys));
|
||
|
|
|
||
|
|
repository
|
||
|
|
.upsert(UpsertUsageRecord {
|
||
|
|
api_key_id: Some("api-key-a".to_string()),
|
||
|
|
status: "completed".to_string(),
|
||
|
|
billing_status: "settled".to_string(),
|
||
|
|
total_tokens: Some(120),
|
||
|
|
total_cost_usd: Some(0.4),
|
||
|
|
created_at_unix_ms: Some(1_711_200_000),
|
||
|
|
updated_at_unix_secs: 1_711_200_000,
|
||
|
|
finalized_at_unix_secs: Some(1_711_200_001),
|
||
|
|
..sample_upsert_usage_record("req-api-move-1")
|
||
|
|
})
|
||
|
|
.await
|
||
|
|
.expect("first upsert should succeed");
|
||
|
|
repository
|
||
|
|
.upsert(UpsertUsageRecord {
|
||
|
|
api_key_id: Some("api-key-b".to_string()),
|
||
|
|
status: "completed".to_string(),
|
||
|
|
billing_status: "settled".to_string(),
|
||
|
|
total_tokens: Some(140),
|
||
|
|
total_cost_usd: Some(0.6),
|
||
|
|
created_at_unix_ms: Some(1_711_200_000),
|
||
|
|
updated_at_unix_secs: 1_711_200_010,
|
||
|
|
finalized_at_unix_secs: Some(1_711_200_011),
|
||
|
|
..sample_upsert_usage_record("req-api-move-1")
|
||
|
|
})
|
||
|
|
.await
|
||
|
|
.expect("moved upsert should succeed");
|
||
|
|
|
||
|
|
let keys = auth_api_keys
|
||
|
|
.list_export_api_keys_by_ids(&["api-key-a".to_string(), "api-key-b".to_string()])
|
||
|
|
.await
|
||
|
|
.expect("key list should succeed");
|
||
|
|
let key_a = keys
|
||
|
|
.iter()
|
||
|
|
.find(|key| key.api_key_id == "api-key-a")
|
||
|
|
.expect("key a should exist");
|
||
|
|
let key_b = keys
|
||
|
|
.iter()
|
||
|
|
.find(|key| key.api_key_id == "api-key-b")
|
||
|
|
.expect("key b should exist");
|
||
|
|
|
||
|
|
assert_eq!(key_a.total_requests, 0);
|
||
|
|
assert_eq!(key_a.total_tokens, 0);
|
||
|
|
assert_eq!(key_a.total_cost_usd, 0.0);
|
||
|
|
|
||
|
|
assert_eq!(key_b.total_requests, 1);
|
||
|
|
assert_eq!(key_b.total_tokens, 140);
|
||
|
|
assert_eq!(key_b.total_cost_usd, 0.6);
|
||
|
|
}
|
||
|
|
|
||
#[tokio::test]
|
|||
|
|
async fn rebuild_provider_key_usage_stats_resets_linked_catalog_to_current_usage() {
|
||
|
|
let provider_catalog = sample_provider_catalog_repository(&["provider-key-1"]);
|
||
|
|
let mut stale_key = provider_catalog
|
||
|
|
.list_keys_by_ids(&["provider-key-1".to_string()])
|
||
|
|
.await
|
||
|
|
.expect("key list should succeed")
|
||
|
|
.into_iter()
|
||
|
|
.next()
|
||
|
|
.expect("provider key should exist");
|
||
|
|
stale_key.request_count = Some(99);
|
||
|
|
stale_key.success_count = Some(88);
|
||
|
|
stale_key.error_count = Some(11);
|
||
|
|
stale_key.total_tokens = 9_999;
|
||
|
|
stale_key.total_cost_usd = 42.0;
|
||
|
|
stale_key.total_response_time_ms = Some(9_999);
|
||
|
|
stale_key.last_used_at_unix_secs = Some(9_999);
|
||
|
|
provider_catalog
|
||
|
|
.update_key(&stale_key)
|
||
|
|
.await
|
||
|
|
.expect("stale key should update");
|
||
|
|
|
||
|
|
let repository = InMemoryUsageReadRepository::seed(vec![
|
||
|
|
sample_usage("req-1", 1_711_300_000),
|
||
|
|
sample_usage("req-2", 1_711_300_250),
|
||
|
|
])
|
||
|
|
.with_provider_catalog_repository(Arc::clone(&provider_catalog));
|
||
|
|
|
||
|
|
let rebuilt = repository
|
||
|
|
.rebuild_provider_api_key_usage_stats()
|
||
|
|
.await
|
||
|
|
.expect("rebuild should succeed");
|
||
|
|
assert_eq!(rebuilt, 1);
|
||
|
|
|
||
|
|
let key = provider_catalog
|
||
|
|
.list_keys_by_ids(&["provider-key-1".to_string()])
|
||
|
|
.await
|
||
|
|
.expect("key list should succeed")
|
||
|
|
.into_iter()
|
||
|
|
.next()
|
||
|
|
.expect("provider key should exist");
|
||
|
|
assert_eq!(key.request_count, Some(2));
|
||
|
|
assert_eq!(key.success_count, Some(2));
|
||
|
|
assert_eq!(key.error_count, Some(0));
|
||
|
|
assert_eq!(key.total_tokens, 300);
|
||
|
|
assert_eq!(key.total_cost_usd, 0.24);
|
||
|
|
assert_eq!(key.total_response_time_ms, Some(840));
|
||
|
|
assert_eq!(key.last_used_at_unix_secs, Some(1_711_300_250));
|
||
|
|
}
|
||
|
|||
|
|
#[tokio::test]
|
||
|
|
async fn rebuild_api_key_usage_stats_resets_linked_auth_export_records_to_current_usage() {
|
||
|
|
let auth_api_keys = sample_auth_api_key_repository(&["api-key-1"]);
|
||
|
|
let mut stale_key = auth_api_keys
|
||
|
|
.list_export_api_keys_by_ids(&["api-key-1".to_string()])
|
||
|
|
.await
|
||
|
|
.expect("key list should succeed")
|
||
|
|
.into_iter()
|
||
|
|
.next()
|
||
|
|
.expect("api key should exist");
|
||
|
|
stale_key.total_requests = 99;
|
||
|
|
stale_key.total_tokens = 9_999;
|
||
|
|
stale_key.total_cost_usd = 42.0;
|
||
|
|
let auth_api_keys = Arc::new(
|
||
|
|
InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||
|
|
Some("hash-api-key-1".to_string()),
|
||
|
|
StoredAuthApiKeySnapshot::new(
|
||
|
|
"user-1".to_string(),
|
||
|
|
"alice".to_string(),
|
||
|
|
Some("[email protected]".to_string()),
|
||
|
|
"user".to_string(),
|
||
|
|
"local".to_string(),
|
||
|
|
true,
|
||
|
|
false,
|
||
|
|
None,
|
||
|
|
None,
|
||
|
|
None,
|
||
|
|
"api-key-1".to_string(),
|
||
|
|
Some("Key api-key-1".to_string()),
|
||
|
|
true,
|
||
|
|
false,
|
||
|
|
false,
|
||
|
|
Some(120),
|
||
|
|
Some(8),
|
||
|
|
None,
|
||
|
|
None,
|
||
|
|
None,
|
||
|
|
None,
|
||
|
|
)
|
||
|
|
.expect("snapshot should build"),
|
||
|
|
)])
|
||
|
|
.with_export_records(vec![stale_key]),
|
||
|
|
);
|
||
|
|
|
||
|
|
let repository = InMemoryUsageReadRepository::seed(vec![
|
||
|
|
sample_usage("req-1", 1_711_300_000),
|
||
|
|
sample_usage("req-2", 1_711_300_250),
|
||
|
|
])
|
||
|
|
.with_auth_api_key_repository(Arc::clone(&auth_api_keys));
|
||
|
|
|
||
|
|
let rebuilt = repository
|
||
|
|
.rebuild_api_key_usage_stats()
|
||
|
|
.await
|
||
|
|
.expect("rebuild should succeed");
|
||
|
|
assert_eq!(rebuilt, 1);
|
||
|
|
|
||
|
|
let key = auth_api_keys
|
||
|
|
.list_export_api_keys_by_ids(&["api-key-1".to_string()])
|
||
|
|
.await
|
||
|
|
.expect("key list should succeed")
|
||
|
|
.into_iter()
|
||
|
|
.next()
|
||
|
|
.expect("api key should exist");
|
||
|
|
assert_eq!(key.total_requests, 2);
|
||
|
|
assert_eq!(key.total_tokens, 300);
|
||
|
|
assert_eq!(key.total_cost_usd, 0.24);
|
||
|
|
}
|
||
|
|||
|
|
#[tokio::test]
|
||
|
|
async fn summarize_usage_provider_performance_computes_tps_and_top_provider_timeline() {
|
||
|
|
let mut first = sample_usage("req-provider-perf-1", 1_711_000_000);
|
||
|
|
first.output_tokens = 60;
|
||
|
|
first.response_time_ms = Some(3000);
|
||
|
|
first.first_byte_time_ms = Some(100);
|
||
|
|
|
||
|
|
let mut second = sample_usage("req-provider-perf-2", 1_711_000_300);
|
||
|
|
second.output_tokens = 40;
|
||
|
|
second.response_time_ms = Some(1000);
|
||
|
|
second.first_byte_time_ms = Some(200);
|
||
second.request_metadata = Some(json!({ "upstream_is_stream": true }));
|
|||
|
|||
|
|
let mut failed = sample_usage("req-provider-perf-failed", 1_711_000_400);
|
||
|
|
failed.output_tokens = 999;
|
||
|
|
failed.response_time_ms = Some(10);
|
||
|
|
failed.first_byte_time_ms = Some(1);
|
||
|
|
failed.status = "failed".to_string();
|
||
|
|
failed.status_code = Some(500);
|
||
|
|
|
||
|
|
let mut other_provider = sample_usage("req-provider-perf-other", 1_711_003_600);
|
||
|
|
other_provider.provider_id = Some("provider-2".to_string());
|
||
|
|
other_provider.provider_name = "Anthropic".to_string();
|
||
|
|
other_provider.output_tokens = 30;
|
||
|
|
other_provider.response_time_ms = Some(3000);
|
||
|
|
other_provider.first_byte_time_ms = None;
|
||
|
|
|
||
|
|
let repository =
|
||
|
|
InMemoryUsageReadRepository::seed(vec![first, second, failed, other_provider]);
|
||
|
|
let summary = repository
|
||
|
|
.summarize_usage_provider_performance(&UsageProviderPerformanceQuery {
|
||
|
|
created_from_unix_secs: 1_711_000_000,
|
||
|
|
created_until_unix_secs: 1_711_010_000,
|
||
|
|
granularity: UsageTimeSeriesGranularity::Hour,
|
||
|
|
tz_offset_minutes: 0,
|
||
|
|
limit: 1,
|
||
provider_id: None,
|
|||
|
|
model: None,
|
||
|
|
api_format: None,
|
||
|
|
endpoint_kind: None,
|
||
|
|
is_stream: None,
|
||
|
|
has_format_conversion: None,
|
||
|
|
slow_threshold_ms: 10_000,
|
||
})
|
|||
|
|
.await
|
||
|
|
.expect("provider performance should summarize");
|
||
|
|
|
||
|
|
assert_eq!(summary.summary.request_count, 4);
|
||
|
|
assert_eq!(summary.summary.success_count, 3);
|
||
assert!((summary.summary.avg_output_tps.expect("summary tps") - 19.117_647).abs() < 0.001);
|
|||
assert_eq!(summary.summary.avg_first_byte_time_ms, Some(150.0));
|
|||
|
|
assert!(
|
||
|
|
(summary
|
||
|
|
.summary
|
||
|
|
.avg_response_time_ms
|
||
|
|
.expect("summary response")
|
||
|
|
- 2333.333)
|
||
|
|
.abs()
|
||
|
|
< 0.001
|
||
|
|
);
|
||
|
|
|
||
|
|
assert_eq!(summary.providers.len(), 1);
|
||
|
|
let provider = &summary.providers[0];
|
||
|
|
assert_eq!(provider.provider_id, "provider-1");
|
||
|
|
assert_eq!(provider.request_count, 3);
|
||
|
|
assert_eq!(provider.success_count, 2);
|
||
|
|
assert_eq!(provider.output_tokens, 1099);
|
||
assert!((provider.avg_output_tps.expect("provider tps") - 26.315_789).abs() < 0.001);
|
|||
assert_eq!(provider.avg_first_byte_time_ms, Some(150.0));
|
|||
|
|
assert_eq!(provider.avg_response_time_ms, Some(2000.0));
|
||
|
|
assert_eq!(provider.p90_response_time_ms, None);
|
||
|
|
assert_eq!(provider.tps_sample_count, 2);
|
||
|
|
assert_eq!(provider.first_byte_sample_count, 2);
|
||
|
|
|
||
|
|
assert_eq!(summary.timeline.len(), 1);
|
||
|
|
assert_eq!(summary.timeline[0].date, "2024-03-21T05:00:00+00:00");
|
||
|
|
assert_eq!(summary.timeline[0].provider_id, "provider-1");
|
||
assert!(
|
|||
|
|
(summary.timeline[0].avg_output_tps.expect("timeline tps") - 26.315_789).abs() < 0.001
|
||
|
|
);
|
||
}
|
|||
}
|