i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-02 09:20:22 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge remote-tracking branch 'origin/aether-rust-pioneer' into aether-rust-pioneer

Browse Source 
# Conflicts:
#	crates/aether-data-contracts/src/repository/usage/mod.rs
#	crates/aether-data/src/repository/global_models/postgres.rs
#	crates/aether-data/src/repository/usage/postgres/mod.rs
This commit is contained in:
fawney19
2026-05-05 18:53:14 +08:00
parent fce7e959e5 825c1b496d
commit aacab1a90c
107 changed files with 7399 additions and 244 deletions
Download Patch File Download Diff File Expand all files Collapse all files

125
crates/aether-data/src/repository/global_models/memory.rs
View File

@@ -497,6 +497,7 @@ impl GlobalModelWriteRepository for InMemoryGlobalModelReadRepository {
Some(global_model.display_name.clone()),
global_model.default_price_per_request,
global_model.default_tiered_pricing.clone(),
global_model.supported_capabilities.clone(),
global_model.config.clone(),
)?;
self.admin_provider_model_items
@@ -542,6 +543,7 @@ impl GlobalModelWriteRepository for InMemoryGlobalModelReadRepository {
existing.global_model_display_name = Some(global_model.display_name.clone());
existing.global_model_default_price_per_request = global_model.default_price_per_request;
existing.global_model_default_tiered_pricing = global_model.default_tiered_pricing.clone();
existing.global_model_supported_capabilities = global_model.supported_capabilities.clone();
existing.global_model_config = global_model.config.clone();
Ok(Some(existing.clone()))
}
@@ -639,8 +641,9 @@ mod tests {
use super::InMemoryGlobalModelReadRepository;
use crate::repository::global_models::{
GlobalModelReadRepository, PublicCatalogModelListQuery, PublicCatalogModelSearchQuery,
PublicGlobalModelQuery, StoredPublicCatalogModel, StoredPublicGlobalModel,
CreateAdminGlobalModelRecord, GlobalModelReadRepository, GlobalModelWriteRepository,
PublicCatalogModelListQuery, PublicCatalogModelSearchQuery, PublicGlobalModelQuery,
StoredPublicCatalogModel, StoredPublicGlobalModel,
};
fn sample_model(
@@ -687,11 +690,83 @@ mod tests {
Some(true),
Some(true),
Some(true),
Some(false),
true,
)
.expect("public catalog model should build")
}
#[tokio::test]
async fn embedding_model_metadata_roundtrip() {
let repository =
InMemoryGlobalModelReadRepository::seed(Vec::<StoredPublicGlobalModel>::new());
let record = CreateAdminGlobalModelRecord::new(
"gm-embedding".to_string(),
"text-embedding-3-small".to_string(),
"Text Embedding 3 Small".to_string(),
true,
None,
Some(json!({"tiers":[{"up_to":null,"input_price_per_1m":0.02}]})),
Some(json!(["embedding"])),
Some(json!({
"api_formats": ["openai:embedding"],
"dimensions": 1536
})),
)
.expect("embedding global model should validate");
repository
.create_admin_global_model(&record)
.await
.expect("embedding global model should persist")
.expect("embedding global model should be returned");
let stored = repository
.get_admin_global_model_by_name("text-embedding-3-small")
.await
.expect("embedding global model should read")
.expect("embedding global model should exist");
assert_eq!(stored.supported_capabilities, Some(json!(["embedding"])));
assert_eq!(
stored
.config
.as_ref()
.and_then(|value| value.get("dimensions")),
Some(&json!(1536))
);
assert_eq!(
stored
.default_tiered_pricing
.as_ref()
.and_then(|value| value.get("tiers"))
.and_then(serde_json::Value::as_array)
.and_then(|tiers| tiers.first())
.and_then(|tier| tier.get("input_price_per_1m"))
.and_then(serde_json::Value::as_f64),
Some(0.02)
);
}
#[tokio::test]
async fn embedding_missing_billing_config_rejected() {
let error = CreateAdminGlobalModelRecord::new(
"gm-embedding".to_string(),
"text-embedding-3-small".to_string(),
"Text Embedding 3 Small".to_string(),
true,
None,
None,
Some(json!(["embedding"])),
None,
)
.expect_err("embedding metadata without billing should fail closed");
assert!(error
.to_string()
.contains("embedding global model requires"));
}
#[tokio::test]
async fn defaults_to_active_models_only() {
let repository = InMemoryGlobalModelReadRepository::seed(vec![
@@ -791,6 +866,52 @@ mod tests {
assert_eq!(items[0].name, "gpt-5");
}
#[tokio::test]
async fn public_catalog_preserves_embedding_capability_without_contaminating_chat_models() {
let mut embedding_model = sample_public_catalog_model(
"model-embedding",
"provider-openai",
"openai",
"text-embedding-3-small",
"text-embedding-3-small",
"Text Embedding 3 Small",
);
embedding_model.supports_embedding = Some(true);
embedding_model.supports_streaming = Some(false);
let chat_model = sample_public_catalog_model(
"model-chat",
"provider-openai",
"openai",
"gpt-5-upstream",
"gpt-5",
"GPT 5",
);
let repository =
InMemoryGlobalModelReadRepository::seed(Vec::<StoredPublicGlobalModel>::new())
.with_public_catalog_models(vec![embedding_model, chat_model]);
let items = repository
.list_public_catalog_models(&PublicCatalogModelListQuery {
provider_id: Some("provider-openai".to_string()),
offset: 0,
limit: 50,
})
.await
.expect("catalog should list");
let embedding = items
.iter()
.find(|item| item.name == "text-embedding-3-small")
.expect("embedding model should be listed");
let chat = items
.iter()
.find(|item| item.name == "gpt-5")
.expect("chat model should be listed");
assert_eq!(embedding.supports_embedding, Some(true));
assert_eq!(embedding.supports_streaming, Some(false));
assert_eq!(chat.supports_embedding, Some(false));
}
#[tokio::test]
async fn searches_public_catalog_models_by_provider_and_display_name() {
let repository =

94
crates/aether-data/src/repository/global_models/mod.rs
View File

@@ -3,6 +3,8 @@ mod mysql;
mod postgres;
mod sqlite;
use serde_json::Value;
#[allow(unused_imports)]
pub(crate) use aether_data_contracts::repository::global_models::{
AdminGlobalModelListQuery, AdminProviderModelListQuery, CreateAdminGlobalModelRecord,
@@ -16,3 +18,95 @@ pub use memory::InMemoryGlobalModelReadRepository;
pub use mysql::MysqlGlobalModelReadRepository;
pub use postgres::SqlxGlobalModelReadRepository;
pub use sqlite::SqliteGlobalModelReadRepository;
const EMBEDDING_CAPABILITY: &str = "embedding";
const EMBEDDING_API_FORMATS: &[&str] = &[
"openai:embedding",
"gemini:embedding",
"jina:embedding",
"doubao:embedding",
"/v1/embeddings",
"/jina/v1/embeddings",
];
pub(super) fn metadata_supports_embedding(
supported_capabilities: Option<&Value>,
global_config: Option<&Value>,
model_config: Option<&Value>,
) -> Option<bool> {
Some(
supported_capabilities.is_some_and(value_contains_embedding_capability)
|| global_config.is_some_and(value_contains_embedding_metadata)
|| model_config.is_some_and(value_contains_embedding_metadata),
)
}
fn value_contains_embedding_capability(value: &Value) -> bool {
match value {
Value::String(value) => value.trim().eq_ignore_ascii_case(EMBEDDING_CAPABILITY),
Value::Array(values) => values.iter().any(value_contains_embedding_capability),
Value::Object(object) => {
object
.get(EMBEDDING_CAPABILITY)
.and_then(Value::as_bool)
.unwrap_or(false)
|| [
"capability",
"model_type",
"type",
"task_type",
"request_type",
]
.iter()
.any(|key| {
object
.get(*key)
.is_some_and(value_contains_embedding_capability)
})
|| ["capabilities", "supported_capabilities"]
.iter()
.any(|key| {
object
.get(*key)
.is_some_and(value_contains_embedding_capability)
})
}
_ => false,
}
}
fn value_contains_embedding_metadata(value: &Value) -> bool {
match value {
Value::String(value) => {
value.trim().eq_ignore_ascii_case(EMBEDDING_CAPABILITY)
|| is_known_embedding_api_format(value)
}
Value::Array(values) => values.iter().any(value_contains_embedding_metadata),
Value::Object(object) => {
value_contains_embedding_capability(value)
|| ["api_format", "client_api_format", "provider_api_format"]
.iter()
.any(|key| {
object
.get(*key)
.and_then(Value::as_str)
.is_some_and(is_known_embedding_api_format)
})
|| ["api_formats", "client_api_formats", "provider_api_formats"]
.iter()
.any(|key| {
object
.get(*key)
.is_some_and(value_contains_embedding_metadata)
})
}
_ => false,
}
}
fn is_known_embedding_api_format(value: &str) -> bool {
let normalized = value.trim().to_ascii_lowercase();
EMBEDDING_API_FORMATS
.iter()
.any(|format| normalized == *format || normalized.ends_with(*format))
}

34
crates/aether-data/src/repository/global_models/mysql.rs
View File

@@ -2,13 +2,13 @@ use async_trait::async_trait;
use sqlx::{mysql::MySqlRow, Row};
use super::{
AdminGlobalModelListQuery, AdminProviderModelListQuery, CreateAdminGlobalModelRecord,
GlobalModelReadRepository, GlobalModelWriteRepository, InMemoryGlobalModelReadRepository,
PublicCatalogModelListQuery, PublicCatalogModelSearchQuery, PublicGlobalModelQuery,
StoredAdminGlobalModel, StoredAdminGlobalModelPage, StoredAdminProviderModel,
StoredProviderActiveGlobalModel, StoredProviderModelStats, StoredPublicCatalogModel,
StoredPublicGlobalModel, StoredPublicGlobalModelPage, UpdateAdminGlobalModelRecord,
UpsertAdminProviderModelRecord,
metadata_supports_embedding, AdminGlobalModelListQuery, AdminProviderModelListQuery,
CreateAdminGlobalModelRecord, GlobalModelReadRepository, GlobalModelWriteRepository,
InMemoryGlobalModelReadRepository, PublicCatalogModelListQuery, PublicCatalogModelSearchQuery,
PublicGlobalModelQuery, StoredAdminGlobalModel, StoredAdminGlobalModelPage,
StoredAdminProviderModel, StoredProviderActiveGlobalModel, StoredProviderModelStats,
StoredPublicCatalogModel, StoredPublicGlobalModel, StoredPublicGlobalModelPage,
UpdateAdminGlobalModelRecord, UpsertAdminProviderModelRecord,
};
use crate::driver::mysql::MysqlPool;
use crate::error::SqlResultExt;
@@ -104,6 +104,7 @@ SELECT
gm.display_name AS global_model_display_name,
gm.default_price_per_request AS global_model_default_price_per_request,
gm.default_tiered_pricing AS global_model_default_tiered_pricing,
gm.supported_capabilities AS global_model_supported_capabilities,
gm.config AS global_model_config
FROM models m
LEFT JOIN global_models gm ON gm.id = m.global_model_id
@@ -130,6 +131,8 @@ SELECT
COALESCE(gm.name, m.provider_model_name) AS name,
COALESCE(NULLIF(gm.display_name, ''), m.provider_model_name) AS display_name,
gm.config AS global_model_config,
gm.supported_capabilities AS global_model_supported_capabilities,
m.config AS model_config,
m.tiered_pricing,
gm.default_tiered_pricing,
m.supports_vision,
@@ -781,6 +784,11 @@ fn map_admin_provider_model_row(
.map_sql_err()?,
"global_models.default_tiered_pricing",
)?,
optional_json_from_string(
row.try_get("global_model_supported_capabilities")
.map_sql_err()?,
"global_models.supported_capabilities",
)?,
optional_json_from_string(
row.try_get("global_model_config").map_sql_err()?,
"global_models.config",
@@ -795,6 +803,13 @@ fn map_public_catalog_model_row(
row.try_get("global_model_config").map_sql_err()?,
"global_models.config",
)?;
let global_model_supported_capabilities = optional_json_from_string(
row.try_get("global_model_supported_capabilities")
.map_sql_err()?,
"global_models.supported_capabilities",
)?;
let model_config =
optional_json_from_string(row.try_get("model_config").map_sql_err()?, "models.config")?;
let tiered_pricing = optional_json_from_string(
row.try_get("tiered_pricing").map_sql_err()?,
"models.tiered_pricing",
@@ -835,6 +850,11 @@ fn map_public_catalog_model_row(
row.try_get("supports_vision").map_sql_err()?,
row.try_get("supports_function_calling").map_sql_err()?,
row.try_get("supports_streaming").map_sql_err()?,
metadata_supports_embedding(
global_model_supported_capabilities.as_ref(),
global_model_config.as_ref(),
model_config.as_ref(),
),
model_is_active && provider_is_active && global_model_is_active,
)
}

65
crates/aether-data/src/repository/global_models/postgres.rs
View File

@@ -60,6 +60,27 @@ SELECT
COALESCE(m.supports_vision, CAST(gm.config->>'vision' AS BOOLEAN), FALSE) AS supports_vision,
COALESCE(m.supports_function_calling, CAST(gm.config->>'function_calling' AS BOOLEAN), FALSE) AS supports_function_calling,
COALESCE(m.supports_streaming, CAST(gm.config->>'streaming' AS BOOLEAN), TRUE) AS supports_streaming,
(
COALESCE(gm.supported_capabilities::jsonb @> '["embedding"]'::jsonb, FALSE)
OR LOWER(COALESCE(gm.config->>'embedding', 'false')) = 'true'
OR LOWER(COALESCE(gm.config->>'model_type', '')) = 'embedding'
OR LOWER(COALESCE(gm.config->>'type', '')) = 'embedding'
OR COALESCE(gm.config->'capabilities' @> '["embedding"]'::jsonb, FALSE)
OR COALESCE(gm.config->'supported_capabilities' @> '["embedding"]'::jsonb, FALSE)
OR COALESCE(gm.config->'api_formats' @> '["openai:embedding"]'::jsonb, FALSE)
OR COALESCE(gm.config->'api_formats' @> '["jina:embedding"]'::jsonb, FALSE)
OR COALESCE(gm.config->'api_formats' @> '["gemini:embedding"]'::jsonb, FALSE)
OR COALESCE(gm.config->'api_formats' @> '["doubao:embedding"]'::jsonb, FALSE)
OR LOWER(COALESCE(m.config->>'embedding', 'false')) = 'true'
OR LOWER(COALESCE(m.config->>'model_type', '')) = 'embedding'
OR LOWER(COALESCE(m.config->>'type', '')) = 'embedding'
OR COALESCE(m.config::jsonb->'capabilities' @> '["embedding"]'::jsonb, FALSE)
OR COALESCE(m.config::jsonb->'supported_capabilities' @> '["embedding"]'::jsonb, FALSE)
OR COALESCE(m.config::jsonb->'api_formats' @> '["openai:embedding"]'::jsonb, FALSE)
OR COALESCE(m.config::jsonb->'api_formats' @> '["jina:embedding"]'::jsonb, FALSE)
OR COALESCE(m.config::jsonb->'api_formats' @> '["gemini:embedding"]'::jsonb, FALSE)
OR COALESCE(m.config::jsonb->'api_formats' @> '["doubao:embedding"]'::jsonb, FALSE)
) AS supports_embedding,
m.is_active
FROM models m
JOIN providers p ON p.id = m.provider_id
@@ -98,6 +119,7 @@ SELECT
gm.display_name AS global_model_display_name,
CAST(gm.default_price_per_request AS DOUBLE PRECISION) AS global_model_default_price_per_request,
gm.default_tiered_pricing AS global_model_default_tiered_pricing,
gm.supported_capabilities AS global_model_supported_capabilities,
gm.config AS global_model_config
FROM models m
LEFT JOIN global_models gm ON gm.id = m.global_model_id
@@ -302,6 +324,7 @@ SELECT
gm.display_name AS global_model_display_name,
CAST(gm.default_price_per_request AS DOUBLE PRECISION) AS global_model_default_price_per_request,
gm.default_tiered_pricing AS global_model_default_tiered_pricing,
gm.supported_capabilities AS global_model_supported_capabilities,
gm.config AS global_model_config
FROM models m
LEFT JOIN global_models gm ON gm.id = m.global_model_id
@@ -348,6 +371,7 @@ SELECT
gm.display_name AS global_model_display_name,
CAST(gm.default_price_per_request AS DOUBLE PRECISION) AS global_model_default_price_per_request,
gm.default_tiered_pricing AS global_model_default_tiered_pricing,
gm.supported_capabilities AS global_model_supported_capabilities,
gm.config AS global_model_config
FROM models m
JOIN global_models gm ON gm.id = m.global_model_id
@@ -487,6 +511,7 @@ SELECT
gm.display_name AS global_model_display_name,
CAST(gm.default_price_per_request AS DOUBLE PRECISION) AS global_model_default_price_per_request,
gm.default_tiered_pricing AS global_model_default_tiered_pricing,
gm.supported_capabilities AS global_model_supported_capabilities,
gm.config AS global_model_config
FROM models m
LEFT JOIN global_models gm ON gm.id = m.global_model_id
@@ -1020,7 +1045,7 @@ fn apply_public_catalog_model_filters(
provider_id: Option<&str>,
search: Option<&str>,
) {
builder.push(" WHERE m.is_active = TRUE AND p.is_active = TRUE");
builder.push(" WHERE m.is_active = TRUE AND COALESCE(m.is_available, TRUE) = TRUE AND p.is_active = TRUE AND COALESCE(gm.is_active, TRUE) = TRUE");
if let Some(provider_id) = provider_id.map(str::trim).filter(|value| !value.is_empty()) {
builder
@@ -1060,6 +1085,7 @@ fn map_public_catalog_model_row(row: &PgRow) -> Result<StoredPublicCatalogModel,
row.try_get("supports_function_calling")
.map_postgres_err()?,
row.try_get("supports_streaming").map_postgres_err()?,
row.try_get("supports_embedding").map_postgres_err()?,
row.try_get("is_active").map_postgres_err()?,
)
}
@@ -1101,6 +1127,8 @@ fn map_admin_provider_model_row(row: &PgRow) -> Result<StoredAdminProviderModel,
.map_postgres_err()?,
row.try_get("global_model_default_tiered_pricing")
.map_postgres_err()?,
row.try_get("global_model_supported_capabilities")
.map_postgres_err()?,
row.try_get("global_model_config").map_postgres_err()?,
)
}
@@ -1177,9 +1205,42 @@ fn map_provider_active_global_model_row(
#[cfg(test)]
mod tests {
use super::SqlxGlobalModelReadRepository;
use super::{SqlxGlobalModelReadRepository, LIST_ADMIN_PROVIDER_MODELS_PREFIX};
use crate::driver::postgres::{PostgresPoolConfig, PostgresPoolFactory};
const ADMIN_PROVIDER_MODEL_REQUIRED_COLUMNS: &[&str] = &[
"global_model_default_tiered_pricing",
"global_model_supported_capabilities",
"global_model_config",
];
fn assert_admin_provider_model_projection_has_required_columns(sql: &str) {
for column in ADMIN_PROVIDER_MODEL_REQUIRED_COLUMNS {
assert!(
sql.contains(column),
"admin provider model SQL projection should include {column}"
);
}
}
#[test]
fn admin_provider_model_sql_projections_include_supported_capabilities() {
assert_admin_provider_model_projection_has_required_columns(
LIST_ADMIN_PROVIDER_MODELS_PREFIX,
);
assert_admin_provider_model_projection_has_required_columns(include_str!("postgres.rs"));
let supported_capabilities_projection = format!(
"{} AS {}",
"gm.supported_capabilities", "global_model_supported_capabilities"
);
assert_eq!(
include_str!("postgres.rs")
.matches(&supported_capabilities_projection)
.count(),
4
);
}
#[tokio::test]
async fn repository_constructs_from_lazy_pool() {
let factory = PostgresPoolFactory::new(PostgresPoolConfig {

34
crates/aether-data/src/repository/global_models/sqlite.rs
View File

@@ -2,13 +2,13 @@ use async_trait::async_trait;
use sqlx::{sqlite::SqliteRow, Row};
use super::{
AdminGlobalModelListQuery, AdminProviderModelListQuery, CreateAdminGlobalModelRecord,
GlobalModelReadRepository, GlobalModelWriteRepository, InMemoryGlobalModelReadRepository,
PublicCatalogModelListQuery, PublicCatalogModelSearchQuery, PublicGlobalModelQuery,
StoredAdminGlobalModel, StoredAdminGlobalModelPage, StoredAdminProviderModel,
StoredProviderActiveGlobalModel, StoredProviderModelStats, StoredPublicCatalogModel,
StoredPublicGlobalModel, StoredPublicGlobalModelPage, UpdateAdminGlobalModelRecord,
UpsertAdminProviderModelRecord,
metadata_supports_embedding, AdminGlobalModelListQuery, AdminProviderModelListQuery,
CreateAdminGlobalModelRecord, GlobalModelReadRepository, GlobalModelWriteRepository,
InMemoryGlobalModelReadRepository, PublicCatalogModelListQuery, PublicCatalogModelSearchQuery,
PublicGlobalModelQuery, StoredAdminGlobalModel, StoredAdminGlobalModelPage,
StoredAdminProviderModel, StoredProviderActiveGlobalModel, StoredProviderModelStats,
StoredPublicCatalogModel, StoredPublicGlobalModel, StoredPublicGlobalModelPage,
UpdateAdminGlobalModelRecord, UpsertAdminProviderModelRecord,
};
use crate::driver::sqlite::SqlitePool;
use crate::error::SqlResultExt;
@@ -117,6 +117,7 @@ SELECT
gm.display_name AS global_model_display_name,
gm.default_price_per_request AS global_model_default_price_per_request,
gm.default_tiered_pricing AS global_model_default_tiered_pricing,
gm.supported_capabilities AS global_model_supported_capabilities,
gm.config AS global_model_config
FROM models m
LEFT JOIN global_models gm ON gm.id = m.global_model_id
@@ -143,6 +144,8 @@ SELECT
COALESCE(gm.name, m.provider_model_name) AS name,
COALESCE(NULLIF(gm.display_name, ''), m.provider_model_name) AS display_name,
gm.config AS global_model_config,
gm.supported_capabilities AS global_model_supported_capabilities,
m.config AS model_config,
m.tiered_pricing,
gm.default_tiered_pricing,
m.supports_vision,
@@ -794,6 +797,11 @@ fn map_admin_provider_model_row(
.map_sql_err()?,
"global_models.default_tiered_pricing",
)?,
optional_json_from_string(
row.try_get("global_model_supported_capabilities")
.map_sql_err()?,
"global_models.supported_capabilities",
)?,
optional_json_from_string(
row.try_get("global_model_config").map_sql_err()?,
"global_models.config",
@@ -808,6 +816,13 @@ fn map_public_catalog_model_row(
row.try_get("global_model_config").map_sql_err()?,
"global_models.config",
)?;
let global_model_supported_capabilities = optional_json_from_string(
row.try_get("global_model_supported_capabilities")
.map_sql_err()?,
"global_models.supported_capabilities",
)?;
let model_config =
optional_json_from_string(row.try_get("model_config").map_sql_err()?, "models.config")?;
let tiered_pricing = optional_json_from_string(
row.try_get("tiered_pricing").map_sql_err()?,
"models.tiered_pricing",
@@ -848,6 +863,11 @@ fn map_public_catalog_model_row(
row.try_get("supports_vision").map_sql_err()?,
row.try_get("supports_function_calling").map_sql_err()?,
row.try_get("supports_streaming").map_sql_err()?,
metadata_supports_embedding(
global_model_supported_capabilities.as_ref(),
global_model_config.as_ref(),
model_config.as_ref(),
),
model_is_active && provider_is_active && global_model_is_active,
)
}

64
crates/aether-data/src/repository/provider_catalog/postgres.rs
View File

@@ -1270,57 +1270,58 @@ INSERT INTO provider_api_keys (
$23,
$24,
$25,
CASE
WHEN $26::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($26::double precision)
END,
$26,
CASE
WHEN $27::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($27::double precision)
END,
$28,
$29,
COALESCE($30, 0),
COALESCE($31, 0),
CASE
WHEN $32::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($32::double precision)
WHEN $28::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($28::double precision)
END,
$29,
$30,
COALESCE($31, 0),
COALESCE($32, 0),
CASE
WHEN $33::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($33::double precision)
END,
$33,
$34,
$35,
$36,
CASE
WHEN $36::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($36::double precision)
WHEN $37::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($37::double precision)
END,
$37,
COALESCE($38, 0),
COALESCE($39, 0),
COALESCE($40, 0),
COALESCE($41, 0),
COALESCE($42, 0),
COALESCE($43, 0),
CASE
WHEN $44::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($44::double precision)
END,
COALESCE($44, 0),
CASE
WHEN $45::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($45::double precision)
END,
$46,
CASE
WHEN $46::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($46::double precision)
END,
$47,
$48,
$49,
CASE
WHEN $50::double precision IS NULL THEN NOW()
ELSE TO_TIMESTAMP($50::double precision)
END,
$50,
CASE
WHEN $51::double precision IS NULL THEN NOW()
ELSE TO_TIMESTAMP($51::double precision)
END,
$52
CASE
WHEN $52::double precision IS NULL THEN NOW()
ELSE TO_TIMESTAMP($52::double precision)
END,
$53
)
"#,
)
@@ -2613,4 +2614,19 @@ mod tests {
assert!(source.contains(".bind(&key.allow_auth_channel_mismatch_formats)"));
assert!(source.contains("row.try_get(\"allow_auth_channel_mismatch_formats\").ok()"));
}
#[test]
fn provider_api_keys_create_key_insert_placeholders_match_bind_order() {
let source = include_str!("postgres.rs");
assert!(source.contains(
" $24,\n $25,\n $26,\n CASE\n WHEN $27::double precision IS NULL THEN NULL"
));
assert!(source.contains(" $29,\n $30,\n COALESCE($31, 0),"));
assert!(source.contains(
" COALESCE($42, 0),\n COALESCE($43, 0),\n COALESCE($44, 0),\n CASE\n WHEN $45::double precision IS NULL THEN NULL"
));
assert!(source.contains(
" CASE\n WHEN $52::double precision IS NULL THEN NOW()\n ELSE TO_TIMESTAMP($52::double precision)\n END,\n $53"
));
}
}

328
crates/aether-data/src/repository/usage/memory.rs
View File

@@ -8,7 +8,9 @@ use aether_data_contracts::repository::usage::{
StoredUsageCacheAffinityIntervalRow, StoredUsageCacheHitSummary, StoredUsageCostSavingsSummary,
StoredUsageDashboardDailyBreakdownRow, StoredUsageDashboardProviderCount,
StoredUsageDashboardSummary, StoredUsageErrorDistributionRow, StoredUsageLeaderboardSummary,
StoredUsagePerformancePercentilesRow, StoredUsageSettledCostSummary,
StoredUsagePerformancePercentilesRow, StoredUsageProviderPerformance,
StoredUsageProviderPerformanceProviderRow, StoredUsageProviderPerformanceSummary,
StoredUsageProviderPerformanceTimelineRow, StoredUsageSettledCostSummary,
StoredUsageTimeSeriesBucket, StoredUsageUserTotals, UsageAuditAggregationGroupBy,
UsageAuditAggregationQuery, UsageAuditKeywordSearchQuery, UsageAuditSummaryQuery,
UsageBodyField, UsageBreakdownGroupBy, UsageBreakdownSummaryQuery,
@@ -17,8 +19,8 @@ use aether_data_contracts::repository::usage::{
UsageDashboardDailyBreakdownQuery, UsageDashboardProviderCountsQuery,
UsageDashboardSummaryQuery, UsageErrorDistributionQuery, UsageLeaderboardGroupBy,
UsageLeaderboardQuery, UsageMonitoringErrorCountQuery, UsageMonitoringErrorListQuery,
UsagePerformancePercentilesQuery, UsageSettledCostSummaryQuery, UsageTimeSeriesGranularity,
UsageTimeSeriesQuery,
UsagePerformancePercentilesQuery, UsageProviderPerformanceQuery, UsageSettledCostSummaryQuery,
UsageTimeSeriesGranularity, UsageTimeSeriesQuery,
};
use async_trait::async_trait;
use chrono::Utc;
@@ -587,6 +589,38 @@ fn usage_matches_performance_percentiles_query(
&& item.status == "completed"
}
fn usage_provider_performance_identity(item: &StoredRequestUsageAudit) -> Option<(String, String)> {
let provider_id = item.provider_id.as_deref()?.trim();
let provider_id_status = provider_id.to_ascii_lowercase();
if provider_id.is_empty() || matches!(provider_id_status.as_str(), "unknown" | "pending") {
return None;
}
let provider_name = item.provider_name.trim();
let provider_name_status = provider_name.to_ascii_lowercase();
if matches!(provider_name_status.as_str(), "unknown" | "pending") {
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;
}
usage_provider_performance_identity(item)
}
fn usage_matches_cost_savings_query(
item: &StoredRequestUsageAudit,
query: &UsageCostSavingsSummaryQuery,
@@ -1650,6 +1684,218 @@ impl UsageReadRepository for InMemoryUsageReadRepository {
.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>,
}
impl ProviderPerformanceBucket {
fn add(&mut self, item: &StoredRequestUsageAudit) {
self.request_count = self.request_count.saturating_add(1);
self.output_tokens = self.output_tokens.saturating_add(item.output_tokens);
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 response_time_ms > 0 && 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(response_time_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);
let bucket = grouped.entry(provider_id).or_default();
if bucket.provider.is_empty() {
bucket.provider = provider;
}
bucket.add(item);
}
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,
),
};
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 p90_first_byte_time_ms =
usage_percentile_cont(&mut bucket.first_byte_times, 0.9);
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,
p90_first_byte_time_ms,
tps_sample_count: bucket.tps_sample_count,
first_byte_sample_count: bucket.first_byte_sample_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);
}
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,
),
},
)
.collect();
Ok(StoredUsageProviderPerformance {
summary,
providers,
timeline,
})
}
async fn summarize_usage_cost_savings(
&self,
query: &UsageCostSavingsSummaryQuery,
@@ -2579,7 +2825,9 @@ mod tests {
StoredProviderUsageWindow, StoredRequestUsageAudit, UpsertUsageRecord, UsageReadRepository,
UsageWriteRepository,
};
use aether_data_contracts::repository::usage::{usage_body_ref, UsageBodyField};
use aether_data_contracts::repository::usage::{
usage_body_ref, UsageBodyField, UsageProviderPerformanceQuery, UsageTimeSeriesGranularity,
};
use serde_json::json;
fn sample_usage(request_id: &str, created_at_unix_ms: i64) -> StoredRequestUsageAudit {
@@ -4562,4 +4810,76 @@ mod tests {
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);
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,
})
.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") - 18.571_428).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_eq!(provider.avg_output_tps, Some(25.0));
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_eq!(summary.timeline[0].avg_output_tps, Some(25.0));
}
}

21
crates/aether-data/src/repository/usage/mod.rs
View File

@@ -241,6 +241,17 @@ macro_rules! impl_materialized_usage_read_repository {
<$crate::repository::usage::InMemoryUsageReadRepository as $crate::repository::usage::UsageReadRepository>::summarize_usage_performance_percentiles(&repository, query).await
}
async fn summarize_usage_provider_performance(
&self,
query: &$crate::repository::usage::UsageProviderPerformanceQuery,
) -> Result<
$crate::repository::usage::StoredUsageProviderPerformance,
$crate::DataLayerError,
> {
let repository = self.materialize_read_model().await?;
<$crate::repository::usage::InMemoryUsageReadRepository as $crate::repository::usage::UsageReadRepository>::summarize_usage_provider_performance(&repository, query).await
}
async fn summarize_usage_cost_savings(
&self,
query: &$crate::repository::usage::UsageCostSavingsSummaryQuery,
@@ -348,7 +359,9 @@ pub(crate) use aether_data_contracts::repository::usage::{
StoredUsageDailySummary, StoredUsageDashboardDailyBreakdownRow,
StoredUsageDashboardProviderCount, StoredUsageDashboardSummary,
StoredUsageErrorDistributionRow, StoredUsageLeaderboardSummary,
StoredUsagePerformancePercentilesRow, StoredUsageSettledCostSummary,
StoredUsagePerformancePercentilesRow, StoredUsageProviderPerformance,
StoredUsageProviderPerformanceProviderRow, StoredUsageProviderPerformanceSummary,
StoredUsageProviderPerformanceTimelineRow, StoredUsageSettledCostSummary,
StoredUsageTimeSeriesBucket, StoredUsageUserTotals, UpsertUsageRecord,
UsageAuditAggregationGroupBy, UsageAuditAggregationQuery, UsageAuditKeywordSearchQuery,
UsageAuditListQuery, UsageAuditSummaryQuery, UsageBreakdownGroupBy, UsageBreakdownSummaryQuery,
@@ -358,9 +371,9 @@ pub(crate) use aether_data_contracts::repository::usage::{
UsageDashboardDailyBreakdownQuery, UsageDashboardProviderCountsQuery,
UsageDashboardSummaryQuery, UsageErrorDistributionQuery, UsageLeaderboardGroupBy,
UsageLeaderboardQuery, UsageMonitoringErrorCountQuery, UsageMonitoringErrorListQuery,
UsagePerformancePercentilesQuery, UsageReadRepository, UsageRepository,
UsageSettledCostSummaryQuery, UsageTimeSeriesGranularity, UsageTimeSeriesQuery,
UsageWriteRepository,
UsagePerformancePercentilesQuery, UsageProviderPerformanceQuery, UsageReadRepository,
UsageRepository, UsageSettledCostSummaryQuery, UsageTimeSeriesGranularity,
UsageTimeSeriesQuery, UsageWriteRepository,
};
pub mod cleanup {
pub use super::postgres::cleanup::*;

416
crates/aether-data/src/repository/usage/postgres/mod.rs
View File

@@ -4,7 +4,9 @@ use aether_data_contracts::repository::usage::{
StoredUsageCacheAffinityIntervalRow, StoredUsageCacheHitSummary, StoredUsageCostSavingsSummary,
StoredUsageDashboardDailyBreakdownRow, StoredUsageDashboardProviderCount,
StoredUsageDashboardSummary, StoredUsageErrorDistributionRow, StoredUsageLeaderboardSummary,
StoredUsagePerformancePercentilesRow, StoredUsageSettledCostSummary,
StoredUsagePerformancePercentilesRow, StoredUsageProviderPerformance,
StoredUsageProviderPerformanceProviderRow, StoredUsageProviderPerformanceSummary,
StoredUsageProviderPerformanceTimelineRow, StoredUsageSettledCostSummary,
StoredUsageTimeSeriesBucket, StoredUsageUserTotals, UsageAuditAggregationGroupBy,
UsageAuditAggregationQuery, UsageAuditKeywordSearchQuery, UsageAuditSummaryQuery,
UsageBodyCaptureState, UsageBodyField, UsageBreakdownGroupBy, UsageBreakdownSummaryQuery,
@@ -13,8 +15,8 @@ use aether_data_contracts::repository::usage::{
UsageCleanupWindow, UsageCostSavingsSummaryQuery, UsageDashboardDailyBreakdownQuery,
UsageDashboardProviderCountsQuery, UsageDashboardSummaryQuery, UsageErrorDistributionQuery,
UsageLeaderboardGroupBy, UsageLeaderboardQuery, UsageMonitoringErrorCountQuery,
UsageMonitoringErrorListQuery, UsagePerformancePercentilesQuery, UsageSettledCostSummaryQuery,
UsageTimeSeriesGranularity, UsageTimeSeriesQuery,
UsageMonitoringErrorListQuery, UsagePerformancePercentilesQuery, UsageProviderPerformanceQuery,
UsageSettledCostSummaryQuery, UsageTimeSeriesGranularity, UsageTimeSeriesQuery,
};
use async_trait::async_trait;
use chrono::{DateTime, Utc};
@@ -806,6 +808,107 @@ fn decode_usage_performance_percentiles_row(
})
}
fn decode_usage_provider_performance_summary(
row: &PgRow,
) -> Result<StoredUsageProviderPerformanceSummary, DataLayerError> {
Ok(StoredUsageProviderPerformanceSummary {
request_count: row
.try_get::<i64, _>("request_count")
.map_postgres_err()?
.max(0) as u64,
success_count: row
.try_get::<i64, _>("success_count")
.map_postgres_err()?
.max(0) as u64,
avg_output_tps: row
.try_get::<Option<f64>, _>("avg_output_tps")
.map_postgres_err()?,
avg_first_byte_time_ms: row
.try_get::<Option<f64>, _>("avg_first_byte_time_ms")
.map_postgres_err()?,
avg_response_time_ms: row
.try_get::<Option<f64>, _>("avg_response_time_ms")
.map_postgres_err()?,
})
}
fn decode_usage_provider_performance_provider_row(
row: &PgRow,
) -> Result<StoredUsageProviderPerformanceProviderRow, DataLayerError> {
Ok(StoredUsageProviderPerformanceProviderRow {
provider_id: row.try_get::<String, _>("provider_id").map_postgres_err()?,
provider: row.try_get::<String, _>("provider").map_postgres_err()?,
request_count: row
.try_get::<i64, _>("request_count")
.map_postgres_err()?
.max(0) as u64,
success_count: row
.try_get::<i64, _>("success_count")
.map_postgres_err()?
.max(0) as u64,
output_tokens: row
.try_get::<i64, _>("output_tokens")
.map_postgres_err()?
.max(0) as u64,
avg_output_tps: row
.try_get::<Option<f64>, _>("avg_output_tps")
.map_postgres_err()?,
avg_first_byte_time_ms: row
.try_get::<Option<f64>, _>("avg_first_byte_time_ms")
.map_postgres_err()?,
avg_response_time_ms: row
.try_get::<Option<f64>, _>("avg_response_time_ms")
.map_postgres_err()?,
p90_response_time_ms: row
.try_get::<Option<i64>, _>("p90_response_time_ms")
.map_postgres_err()?
.map(|value| value.max(0) as u64),
p90_first_byte_time_ms: row
.try_get::<Option<i64>, _>("p90_first_byte_time_ms")
.map_postgres_err()?
.map(|value| value.max(0) as u64),
tps_sample_count: row
.try_get::<i64, _>("tps_sample_count")
.map_postgres_err()?
.max(0) as u64,
first_byte_sample_count: row
.try_get::<i64, _>("first_byte_sample_count")
.map_postgres_err()?
.max(0) as u64,
})
}
fn decode_usage_provider_performance_timeline_row(
row: &PgRow,
) -> Result<StoredUsageProviderPerformanceTimelineRow, DataLayerError> {
Ok(StoredUsageProviderPerformanceTimelineRow {
date: row.try_get::<String, _>("date").map_postgres_err()?,
provider_id: row.try_get::<String, _>("provider_id").map_postgres_err()?,
provider: row.try_get::<String, _>("provider").map_postgres_err()?,
request_count: row
.try_get::<i64, _>("request_count")
.map_postgres_err()?
.max(0) as u64,
success_count: row
.try_get::<i64, _>("success_count")
.map_postgres_err()?
.max(0) as u64,
output_tokens: row
.try_get::<i64, _>("output_tokens")
.map_postgres_err()?
.max(0) as u64,
avg_output_tps: row
.try_get::<Option<f64>, _>("avg_output_tps")
.map_postgres_err()?,
avg_first_byte_time_ms: row
.try_get::<Option<f64>, _>("avg_first_byte_time_ms")
.map_postgres_err()?,
avg_response_time_ms: row
.try_get::<Option<f64>, _>("avg_response_time_ms")
.map_postgres_err()?,
})
}
fn decode_usage_time_series_bucket_row(
row: &PgRow,
) -> Result<StoredUsageTimeSeriesBucket, DataLayerError> {
@@ -4467,6 +4570,306 @@ ORDER BY date ASC
Ok(items)
}
async fn summarize_usage_provider_performance_summary(
&self,
query: &UsageProviderPerformanceQuery,
) -> Result<StoredUsageProviderPerformanceSummary, DataLayerError> {
let mut builder = QueryBuilder::<Postgres>::new(
r#"
WITH filtered_usage AS (
SELECT
GREATEST(COALESCE("usage".output_tokens, 0), 0) AS output_tokens,
GREATEST(COALESCE("usage".response_time_ms, 0), 0) AS response_time_ms,
GREATEST(COALESCE("usage".first_byte_time_ms, 0), 0) AS first_byte_time_ms,
"usage".response_time_ms IS NOT NULL AS has_response_time,
"usage".first_byte_time_ms IS NOT NULL AS has_first_byte_time,
CASE
WHEN lower(COALESCE("usage".status, '')) IN ('completed', 'success', 'ok', 'billed', 'settled')
AND ("usage".status_code IS NULL OR "usage".status_code < 400)
THEN 1
ELSE 0
END AS success_flag
FROM usage_billing_facts AS "usage"
WHERE "usage".created_at >= TO_TIMESTAMP("#,
);
builder.push_bind(query.created_from_unix_secs as f64);
builder.push(
r#"::double precision)
AND "usage".created_at < TO_TIMESTAMP("#,
);
builder.push_bind(query.created_until_unix_secs as f64);
builder.push(
r#"::double precision)
AND COALESCE("usage".status, '') NOT IN ('pending', 'streaming')
AND NULLIF(BTRIM(COALESCE("usage".provider_id, '')), '') IS NOT NULL
AND lower(BTRIM(COALESCE("usage".provider_id, ''))) NOT IN ('unknown', 'pending')
AND lower(BTRIM(COALESCE("usage".provider_name, ''))) NOT IN ('unknown', 'pending')
)
SELECT
COUNT(*)::BIGINT AS request_count,
COALESCE(SUM(success_flag), 0)::BIGINT AS success_count,
CASE
WHEN COALESCE(SUM(CASE
WHEN success_flag = 1 AND response_time_ms > 0 AND output_tokens > 0
THEN response_time_ms
ELSE 0
END), 0) > 0
THEN COALESCE(SUM(CASE
WHEN success_flag = 1 AND response_time_ms > 0 AND output_tokens > 0
THEN output_tokens
ELSE 0
END), 0)::DOUBLE PRECISION * 1000.0 / COALESCE(SUM(CASE
WHEN success_flag = 1 AND response_time_ms > 0 AND output_tokens > 0
THEN response_time_ms
ELSE 0
END), 0)::DOUBLE PRECISION
ELSE NULL
END AS avg_output_tps,
AVG(first_byte_time_ms::DOUBLE PRECISION)
FILTER (WHERE success_flag = 1 AND has_first_byte_time) AS avg_first_byte_time_ms,
AVG(response_time_ms::DOUBLE PRECISION)
FILTER (WHERE success_flag = 1 AND has_response_time) AS avg_response_time_ms
FROM filtered_usage
"#,
);
let row = builder
.build()
.fetch_one(&self.pool)
.await
.map_postgres_err()?;
decode_usage_provider_performance_summary(&row)
}
async fn summarize_usage_provider_performance_providers(
&self,
query: &UsageProviderPerformanceQuery,
) -> Result<Vec<StoredUsageProviderPerformanceProviderRow>, DataLayerError> {
let mut builder = QueryBuilder::<Postgres>::new(
r#"
WITH filtered_usage AS (
SELECT
COALESCE("usage".provider_id, '') AS provider_id,
COALESCE(NULLIF(BTRIM("usage".provider_name), ''), COALESCE("usage".provider_id, '')) AS provider,
GREATEST(COALESCE("usage".output_tokens, 0), 0) AS output_tokens,
GREATEST(COALESCE("usage".response_time_ms, 0), 0) AS response_time_ms,
GREATEST(COALESCE("usage".first_byte_time_ms, 0), 0) AS first_byte_time_ms,
"usage".response_time_ms IS NOT NULL AS has_response_time,
"usage".first_byte_time_ms IS NOT NULL AS has_first_byte_time,
CASE
WHEN lower(COALESCE("usage".status, '')) IN ('completed', 'success', 'ok', 'billed', 'settled')
AND ("usage".status_code IS NULL OR "usage".status_code < 400)
THEN 1
ELSE 0
END AS success_flag
FROM usage_billing_facts AS "usage"
WHERE "usage".created_at >= TO_TIMESTAMP("#,
);
builder.push_bind(query.created_from_unix_secs as f64);
builder.push(
r#"::double precision)
AND "usage".created_at < TO_TIMESTAMP("#,
);
builder.push_bind(query.created_until_unix_secs as f64);
builder.push(
r#"::double precision)
AND COALESCE("usage".status, '') NOT IN ('pending', 'streaming')
AND NULLIF(BTRIM(COALESCE("usage".provider_id, '')), '') IS NOT NULL
AND lower(BTRIM(COALESCE("usage".provider_id, ''))) NOT IN ('unknown', 'pending')
AND lower(BTRIM(COALESCE("usage".provider_name, ''))) NOT IN ('unknown', 'pending')
)
SELECT
provider_id,
COALESCE(MAX(NULLIF(provider, '')), provider_id) AS provider,
COUNT(*)::BIGINT AS request_count,
COALESCE(SUM(success_flag), 0)::BIGINT AS success_count,
COALESCE(SUM(output_tokens), 0)::BIGINT AS output_tokens,
CASE
WHEN COALESCE(SUM(CASE
WHEN success_flag = 1 AND response_time_ms > 0 AND output_tokens > 0
THEN response_time_ms
ELSE 0
END), 0) > 0
THEN COALESCE(SUM(CASE
WHEN success_flag = 1 AND response_time_ms > 0 AND output_tokens > 0
THEN output_tokens
ELSE 0
END), 0)::DOUBLE PRECISION * 1000.0 / COALESCE(SUM(CASE
WHEN success_flag = 1 AND response_time_ms > 0 AND output_tokens > 0
THEN response_time_ms
ELSE 0
END), 0)::DOUBLE PRECISION
ELSE NULL
END AS avg_output_tps,
AVG(first_byte_time_ms::DOUBLE PRECISION)
FILTER (WHERE success_flag = 1 AND has_first_byte_time) AS avg_first_byte_time_ms,
AVG(response_time_ms::DOUBLE PRECISION)
FILTER (WHERE success_flag = 1 AND has_response_time) AS avg_response_time_ms,
CASE
WHEN COUNT(response_time_ms) FILTER (WHERE success_flag = 1 AND has_response_time) >= 10
THEN FLOOR(PERCENTILE_CONT(0.9) WITHIN GROUP (ORDER BY response_time_ms)
FILTER (WHERE success_flag = 1 AND has_response_time))::BIGINT
ELSE NULL
END AS p90_response_time_ms,
CASE
WHEN COUNT(first_byte_time_ms) FILTER (WHERE success_flag = 1 AND has_first_byte_time) >= 10
THEN FLOOR(PERCENTILE_CONT(0.9) WITHIN GROUP (ORDER BY first_byte_time_ms)
FILTER (WHERE success_flag = 1 AND has_first_byte_time))::BIGINT
ELSE NULL
END AS p90_first_byte_time_ms,
COALESCE(SUM(CASE
WHEN success_flag = 1 AND response_time_ms > 0 AND output_tokens > 0
THEN 1
ELSE 0
END), 0)::BIGINT AS tps_sample_count,
(COUNT(first_byte_time_ms) FILTER (WHERE success_flag = 1 AND has_first_byte_time))::BIGINT
AS first_byte_sample_count
FROM filtered_usage
GROUP BY provider_id
ORDER BY request_count DESC, provider_id ASC
"#,
);
let mut rows = builder.build().fetch(&self.pool);
let mut items = Vec::new();
while let Some(row) = rows.try_next().await.map_postgres_err()? {
items.push(decode_usage_provider_performance_provider_row(&row)?);
}
Ok(items)
}
async fn summarize_usage_provider_performance_timeline(
&self,
query: &UsageProviderPerformanceQuery,
provider_ids: &[String],
) -> Result<Vec<StoredUsageProviderPerformanceTimelineRow>, DataLayerError> {
if provider_ids.is_empty() {
return Ok(Vec::new());
}
let mut builder = QueryBuilder::<Postgres>::new("WITH filtered_usage AS ( SELECT ");
match query.granularity {
UsageTimeSeriesGranularity::Day => {
builder
.push("TO_CHAR(date_trunc('day', \"usage\".created_at + (")
.push_bind(query.tz_offset_minutes)
.push("::integer * INTERVAL '1 minute')), 'YYYY-MM-DD') AS date");
}
UsageTimeSeriesGranularity::Hour => {
builder
.push("TO_CHAR(date_trunc('hour', \"usage\".created_at + (")
.push_bind(query.tz_offset_minutes)
.push("::integer * INTERVAL '1 minute')), 'YYYY-MM-DD\"T\"HH24:00:00+00:00') AS date");
}
}
builder.push(
r#",
COALESCE("usage".provider_id, '') AS provider_id,
COALESCE(NULLIF(BTRIM("usage".provider_name), ''), COALESCE("usage".provider_id, '')) AS provider,
GREATEST(COALESCE("usage".output_tokens, 0), 0) AS output_tokens,
GREATEST(COALESCE("usage".response_time_ms, 0), 0) AS response_time_ms,
GREATEST(COALESCE("usage".first_byte_time_ms, 0), 0) AS first_byte_time_ms,
"usage".response_time_ms IS NOT NULL AS has_response_time,
"usage".first_byte_time_ms IS NOT NULL AS has_first_byte_time,
CASE
WHEN lower(COALESCE("usage".status, '')) IN ('completed', 'success', 'ok', 'billed', 'settled')
AND ("usage".status_code IS NULL OR "usage".status_code < 400)
THEN 1
ELSE 0
END AS success_flag
FROM usage_billing_facts AS "usage"
WHERE "usage".created_at >= TO_TIMESTAMP("#,
);
builder.push_bind(query.created_from_unix_secs as f64);
builder.push(
r#"::double precision)
AND "usage".created_at < TO_TIMESTAMP("#,
);
builder.push_bind(query.created_until_unix_secs as f64);
builder.push(
r#"::double precision)
AND COALESCE("usage".status, '') NOT IN ('pending', 'streaming')
AND NULLIF(BTRIM(COALESCE("usage".provider_id, '')), '') IS NOT NULL
AND lower(BTRIM(COALESCE("usage".provider_id, ''))) NOT IN ('unknown', 'pending')
AND lower(BTRIM(COALESCE("usage".provider_name, ''))) NOT IN ('unknown', 'pending')
AND "usage".provider_id = ANY("#,
);
builder.push_bind(provider_ids.to_vec());
builder.push(
r#")
)
SELECT
date,
provider_id,
COALESCE(MAX(NULLIF(provider, '')), provider_id) AS provider,
COUNT(*)::BIGINT AS request_count,
COALESCE(SUM(success_flag), 0)::BIGINT AS success_count,
COALESCE(SUM(output_tokens), 0)::BIGINT AS output_tokens,
CASE
WHEN COALESCE(SUM(CASE
WHEN success_flag = 1 AND response_time_ms > 0 AND output_tokens > 0
THEN response_time_ms
ELSE 0
END), 0) > 0
THEN COALESCE(SUM(CASE
WHEN success_flag = 1 AND response_time_ms > 0 AND output_tokens > 0
THEN output_tokens
ELSE 0
END), 0)::DOUBLE PRECISION * 1000.0 / COALESCE(SUM(CASE
WHEN success_flag = 1 AND response_time_ms > 0 AND output_tokens > 0
THEN response_time_ms
ELSE 0
END), 0)::DOUBLE PRECISION
ELSE NULL
END AS avg_output_tps,
AVG(first_byte_time_ms::DOUBLE PRECISION)
FILTER (WHERE success_flag = 1 AND has_first_byte_time) AS avg_first_byte_time_ms,
AVG(response_time_ms::DOUBLE PRECISION)
FILTER (WHERE success_flag = 1 AND has_response_time) AS avg_response_time_ms
FROM filtered_usage
GROUP BY date, provider_id
ORDER BY date ASC, provider_id ASC
"#,
);
let mut rows = builder.build().fetch(&self.pool);
let mut items = Vec::new();
while let Some(row) = rows.try_next().await.map_postgres_err()? {
items.push(decode_usage_provider_performance_timeline_row(&row)?);
}
Ok(items)
}
pub async fn summarize_usage_provider_performance(
&self,
query: &UsageProviderPerformanceQuery,
) -> Result<StoredUsageProviderPerformance, DataLayerError> {
if query.created_from_unix_secs >= query.created_until_unix_secs {
return Ok(StoredUsageProviderPerformance::default());
}
let summary = self
.summarize_usage_provider_performance_summary(query)
.await?;
let mut providers = self
.summarize_usage_provider_performance_providers(query)
.await?;
providers.truncate(query.limit.max(1));
let provider_ids = providers
.iter()
.map(|row| row.provider_id.clone())
.collect::<Vec<_>>();
let timeline = self
.summarize_usage_provider_performance_timeline(query, &provider_ids)
.await?;
Ok(StoredUsageProviderPerformance {
summary,
providers,
timeline,
})
}
async fn summarize_usage_cost_savings_raw_from_range(
&self,
start_utc: DateTime<Utc>,
@@ -6999,6 +7402,13 @@ impl UsageReadRepository for SqlxUsageReadRepository {
Self::summarize_usage_performance_percentiles(self, query).await
}
async fn summarize_usage_provider_performance(
&self,
query: &UsageProviderPerformanceQuery,
) -> Result<StoredUsageProviderPerformance, DataLayerError> {
Self::summarize_usage_provider_performance(self, query).await
}
async fn summarize_usage_cost_savings(
&self,
query: &UsageCostSavingsSummaryQuery,