mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-06 01:17:46 +08:00
1349 lines
45 KiB
Rust
1349 lines
45 KiB
Rust
use async_trait::async_trait;
|
|
use futures_util::{stream::TryStream, TryStreamExt};
|
|
use serde_json::Value;
|
|
use sqlx::{postgres::PgRow, PgPool, Postgres, QueryBuilder, Row};
|
|
|
|
use aether_data_contracts::repository::global_models::{
|
|
AdminGlobalModelListQuery, AdminProviderModelListQuery, CreateAdminGlobalModelRecord,
|
|
GlobalModelReadRepository, GlobalModelWriteRepository, PublicCatalogModelListQuery,
|
|
PublicCatalogModelSearchQuery, PublicGlobalModelQuery, StoredAdminGlobalModel,
|
|
StoredAdminGlobalModelPage, StoredAdminProviderModel, StoredProviderActiveGlobalModel,
|
|
StoredProviderModelStats, StoredPublicCatalogModel, StoredPublicGlobalModel,
|
|
StoredPublicGlobalModelPage, UpdateAdminGlobalModelRecord, UpsertAdminProviderModelRecord,
|
|
};
|
|
use aether_data_contracts::DataLayerError;
|
|
|
|
use crate::error::SqlxResultExt;
|
|
|
|
const LIST_PUBLIC_GLOBAL_MODELS_PREFIX: &str = r#"
|
|
SELECT
|
|
id,
|
|
name,
|
|
display_name,
|
|
is_active,
|
|
CAST(default_price_per_request AS DOUBLE PRECISION) AS default_price_per_request,
|
|
default_tiered_pricing,
|
|
supported_capabilities,
|
|
config
|
|
FROM global_models
|
|
"#;
|
|
|
|
const COUNT_PUBLIC_GLOBAL_MODELS_PREFIX: &str = r#"
|
|
SELECT COUNT(id) AS total
|
|
FROM global_models
|
|
"#;
|
|
|
|
const LIST_PUBLIC_CATALOG_MODELS_PREFIX: &str = r#"
|
|
SELECT
|
|
m.id,
|
|
m.provider_id,
|
|
p.name AS provider_name,
|
|
m.provider_model_name,
|
|
COALESCE(gm.name, m.provider_model_name) AS name,
|
|
COALESCE(NULLIF(gm.display_name, ''), m.provider_model_name) AS display_name,
|
|
CASE
|
|
WHEN gm.config IS NULL THEN NULL
|
|
ELSE gm.config->>'description'
|
|
END AS description,
|
|
CASE
|
|
WHEN gm.config IS NULL THEN NULL
|
|
ELSE gm.config->>'icon_url'
|
|
END AS icon_url,
|
|
COALESCE(
|
|
CAST((COALESCE(m.tiered_pricing, gm.default_tiered_pricing)->'tiers'->0->>'input_price_per_1m') AS DOUBLE PRECISION),
|
|
0.0
|
|
) AS input_price_per_1m,
|
|
COALESCE(
|
|
CAST((COALESCE(m.tiered_pricing, gm.default_tiered_pricing)->'tiers'->0->>'output_price_per_1m') AS DOUBLE PRECISION),
|
|
0.0
|
|
) AS output_price_per_1m,
|
|
CAST((COALESCE(m.tiered_pricing, gm.default_tiered_pricing)->'tiers'->0->>'cache_creation_price_per_1m') AS DOUBLE PRECISION) AS cache_creation_price_per_1m,
|
|
CAST((COALESCE(m.tiered_pricing, gm.default_tiered_pricing)->'tiers'->0->>'cache_read_price_per_1m') AS DOUBLE PRECISION) AS cache_read_price_per_1m,
|
|
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 COALESCE(gm.config->'api_formats' @> '["aliyun:multimodal_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)
|
|
OR COALESCE(m.config::jsonb->'api_formats' @> '["aliyun:multimodal_embedding"]'::jsonb, FALSE)
|
|
) AS supports_embedding,
|
|
m.is_active
|
|
FROM models m
|
|
JOIN providers p ON p.id = m.provider_id
|
|
LEFT JOIN global_models gm ON gm.id = m.global_model_id
|
|
"#;
|
|
|
|
const LIST_PROVIDER_MODEL_STATS_PREFIX: &str = r#"
|
|
SELECT
|
|
provider_id,
|
|
COUNT(id) AS total_models,
|
|
COALESCE(SUM(CASE WHEN is_active = TRUE THEN 1 ELSE 0 END), 0) AS active_models
|
|
FROM models
|
|
WHERE provider_id IN (
|
|
"#;
|
|
|
|
const LIST_ADMIN_PROVIDER_MODELS_PREFIX: &str = r#"
|
|
SELECT
|
|
m.id,
|
|
m.provider_id,
|
|
m.global_model_id,
|
|
m.provider_model_name,
|
|
m.provider_model_mappings,
|
|
CAST(m.price_per_request AS DOUBLE PRECISION) AS price_per_request,
|
|
m.tiered_pricing,
|
|
m.supports_vision,
|
|
m.supports_function_calling,
|
|
m.supports_streaming,
|
|
m.supports_extended_thinking,
|
|
m.supports_image_generation,
|
|
m.is_active,
|
|
COALESCE(m.is_available, TRUE) AS is_available,
|
|
m.config,
|
|
EXTRACT(EPOCH FROM m.created_at)::bigint AS created_at_unix_ms,
|
|
EXTRACT(EPOCH FROM m.updated_at)::bigint AS updated_at_unix_secs,
|
|
gm.name AS global_model_name,
|
|
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
|
|
"#;
|
|
|
|
const LIST_ADMIN_GLOBAL_MODELS_PREFIX: &str = r#"
|
|
SELECT
|
|
gm.id,
|
|
gm.name,
|
|
COALESCE(NULLIF(gm.display_name, ''), gm.name) AS display_name,
|
|
gm.is_active,
|
|
CAST(gm.default_price_per_request AS DOUBLE PRECISION) AS default_price_per_request,
|
|
gm.default_tiered_pricing,
|
|
gm.supported_capabilities,
|
|
gm.config,
|
|
COALESCE(gm_stats.provider_count, 0) AS provider_count,
|
|
COALESCE(gm_stats.active_provider_count, 0) AS active_provider_count,
|
|
COALESCE(gm.usage_count, 0)::bigint AS usage_count,
|
|
EXTRACT(EPOCH FROM gm.created_at)::bigint AS created_at_unix_ms,
|
|
EXTRACT(EPOCH FROM gm.updated_at)::bigint AS updated_at_unix_secs
|
|
FROM global_models gm
|
|
LEFT JOIN (
|
|
SELECT
|
|
m.global_model_id,
|
|
COUNT(DISTINCT m.provider_id)::bigint AS provider_count,
|
|
COUNT(
|
|
DISTINCT CASE
|
|
WHEN m.is_active = TRUE AND COALESCE(m.is_available, TRUE) = TRUE AND p.is_active = TRUE THEN m.provider_id
|
|
ELSE NULL
|
|
END
|
|
)::bigint AS active_provider_count
|
|
FROM models m
|
|
JOIN providers p ON p.id = m.provider_id
|
|
GROUP BY m.global_model_id
|
|
) gm_stats ON gm_stats.global_model_id = gm.id
|
|
"#;
|
|
|
|
const COUNT_ADMIN_GLOBAL_MODELS_PREFIX: &str = r#"
|
|
SELECT COUNT(id) AS total
|
|
FROM global_models gm
|
|
"#;
|
|
|
|
const LIST_ACTIVE_GLOBAL_MODEL_IDS_BY_PROVIDER_IDS_PREFIX: &str = r#"
|
|
SELECT DISTINCT
|
|
provider_id,
|
|
global_model_id
|
|
FROM models
|
|
WHERE provider_id IN (
|
|
"#;
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub struct SqlxGlobalModelReadRepository {
|
|
pool: PgPool,
|
|
}
|
|
|
|
impl SqlxGlobalModelReadRepository {
|
|
pub fn new(pool: PgPool) -> Self {
|
|
Self { pool }
|
|
}
|
|
|
|
pub async fn list_public_models(
|
|
&self,
|
|
query: &PublicGlobalModelQuery,
|
|
) -> Result<StoredPublicGlobalModelPage, DataLayerError> {
|
|
let mut count_builder = QueryBuilder::<Postgres>::new(COUNT_PUBLIC_GLOBAL_MODELS_PREFIX);
|
|
apply_public_model_filters(&mut count_builder, query);
|
|
let count_row = count_builder
|
|
.build()
|
|
.fetch_one(&self.pool)
|
|
.await
|
|
.map_postgres_err()?;
|
|
let total = count_row
|
|
.try_get::<i64, _>("total")
|
|
.map(|value| value.max(0) as usize)
|
|
.map_postgres_err()?;
|
|
|
|
let mut list_builder = QueryBuilder::<Postgres>::new(LIST_PUBLIC_GLOBAL_MODELS_PREFIX);
|
|
apply_public_model_filters(&mut list_builder, query);
|
|
list_builder
|
|
.push(" ORDER BY name ASC OFFSET ")
|
|
.push_bind(query.offset as i64)
|
|
.push(" LIMIT ")
|
|
.push_bind(query.limit as i64);
|
|
let query = list_builder.build();
|
|
let items = collect_query_rows(query.fetch(&self.pool), map_row).await?;
|
|
|
|
Ok(StoredPublicGlobalModelPage { items, total })
|
|
}
|
|
|
|
pub async fn list_provider_model_stats(
|
|
&self,
|
|
provider_ids: &[String],
|
|
) -> Result<Vec<StoredProviderModelStats>, DataLayerError> {
|
|
if provider_ids.is_empty() {
|
|
return Ok(Vec::new());
|
|
}
|
|
|
|
let mut builder = build_provider_id_list_query(
|
|
LIST_PROVIDER_MODEL_STATS_PREFIX,
|
|
provider_ids,
|
|
")\nGROUP BY provider_id\nORDER BY provider_id ASC",
|
|
);
|
|
let query = builder.build();
|
|
collect_query_rows(query.fetch(&self.pool), map_provider_model_stats_row).await
|
|
}
|
|
|
|
pub async fn list_active_global_model_ids_by_provider_ids(
|
|
&self,
|
|
provider_ids: &[String],
|
|
) -> Result<Vec<StoredProviderActiveGlobalModel>, DataLayerError> {
|
|
if provider_ids.is_empty() {
|
|
return Ok(Vec::new());
|
|
}
|
|
|
|
let mut builder = build_provider_id_list_query(
|
|
LIST_ACTIVE_GLOBAL_MODEL_IDS_BY_PROVIDER_IDS_PREFIX,
|
|
provider_ids,
|
|
")\nAND is_active = TRUE\nAND global_model_id IS NOT NULL\nORDER BY provider_id ASC, global_model_id ASC",
|
|
);
|
|
let query = builder.build();
|
|
collect_query_rows(
|
|
query.fetch(&self.pool),
|
|
map_provider_active_global_model_row,
|
|
)
|
|
.await
|
|
}
|
|
|
|
pub async fn list_admin_provider_models(
|
|
&self,
|
|
query: &AdminProviderModelListQuery,
|
|
) -> Result<Vec<StoredAdminProviderModel>, DataLayerError> {
|
|
let mut builder = QueryBuilder::<Postgres>::new(LIST_ADMIN_PROVIDER_MODELS_PREFIX);
|
|
builder
|
|
.push(" WHERE m.provider_id = ")
|
|
.push_bind(query.provider_id.trim().to_string());
|
|
if let Some(is_active) = query.is_active {
|
|
builder.push(" AND m.is_active = ").push_bind(is_active);
|
|
}
|
|
builder
|
|
.push(" ORDER BY m.created_at DESC, m.id ASC OFFSET ")
|
|
.push_bind(query.offset as i64)
|
|
.push(" LIMIT ")
|
|
.push_bind(query.limit as i64);
|
|
let query = builder.build();
|
|
collect_query_rows(query.fetch(&self.pool), map_admin_provider_model_row).await
|
|
}
|
|
|
|
pub async fn list_admin_global_models(
|
|
&self,
|
|
query: &AdminGlobalModelListQuery,
|
|
) -> Result<StoredAdminGlobalModelPage, DataLayerError> {
|
|
let mut count_builder = QueryBuilder::<Postgres>::new(COUNT_ADMIN_GLOBAL_MODELS_PREFIX);
|
|
apply_admin_global_model_filters(&mut count_builder, query);
|
|
let count_row = count_builder
|
|
.build()
|
|
.fetch_one(&self.pool)
|
|
.await
|
|
.map_postgres_err()?;
|
|
let total = count_row
|
|
.try_get::<i64, _>("total")
|
|
.map(|value| value.max(0) as usize)
|
|
.map_postgres_err()?;
|
|
|
|
let mut list_builder = QueryBuilder::<Postgres>::new(LIST_ADMIN_GLOBAL_MODELS_PREFIX);
|
|
apply_admin_global_model_filters(&mut list_builder, query);
|
|
list_builder
|
|
.push(" ORDER BY name ASC OFFSET ")
|
|
.push_bind(query.offset as i64)
|
|
.push(" LIMIT ")
|
|
.push_bind(query.limit as i64);
|
|
let query = list_builder.build();
|
|
let items = collect_query_rows(query.fetch(&self.pool), map_admin_global_model_row).await?;
|
|
Ok(StoredAdminGlobalModelPage { items, total })
|
|
}
|
|
|
|
pub async fn get_admin_provider_model(
|
|
&self,
|
|
provider_id: &str,
|
|
model_id: &str,
|
|
) -> Result<Option<StoredAdminProviderModel>, DataLayerError> {
|
|
let row = sqlx::query(
|
|
r#"
|
|
SELECT
|
|
m.id,
|
|
m.provider_id,
|
|
m.global_model_id,
|
|
m.provider_model_name,
|
|
m.provider_model_mappings,
|
|
CAST(m.price_per_request AS DOUBLE PRECISION) AS price_per_request,
|
|
m.tiered_pricing,
|
|
m.supports_vision,
|
|
m.supports_function_calling,
|
|
m.supports_streaming,
|
|
m.supports_extended_thinking,
|
|
m.supports_image_generation,
|
|
m.is_active,
|
|
COALESCE(m.is_available, TRUE) AS is_available,
|
|
m.config,
|
|
EXTRACT(EPOCH FROM m.created_at)::bigint AS created_at_unix_ms,
|
|
EXTRACT(EPOCH FROM m.updated_at)::bigint AS updated_at_unix_secs,
|
|
gm.name AS global_model_name,
|
|
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
|
|
WHERE m.provider_id = $1
|
|
AND m.id = $2
|
|
LIMIT 1
|
|
"#,
|
|
)
|
|
.bind(provider_id)
|
|
.bind(model_id)
|
|
.fetch_optional(&self.pool)
|
|
.await
|
|
.map_postgres_err()?;
|
|
|
|
row.as_ref().map(map_admin_provider_model_row).transpose()
|
|
}
|
|
|
|
pub async fn list_admin_provider_available_source_models(
|
|
&self,
|
|
provider_id: &str,
|
|
) -> Result<Vec<StoredAdminProviderModel>, DataLayerError> {
|
|
collect_query_rows(
|
|
sqlx::query(
|
|
r#"
|
|
SELECT
|
|
m.id,
|
|
m.provider_id,
|
|
m.global_model_id,
|
|
m.provider_model_name,
|
|
m.provider_model_mappings,
|
|
CAST(m.price_per_request AS DOUBLE PRECISION) AS price_per_request,
|
|
m.tiered_pricing,
|
|
m.supports_vision,
|
|
m.supports_function_calling,
|
|
m.supports_streaming,
|
|
m.supports_extended_thinking,
|
|
m.supports_image_generation,
|
|
m.is_active,
|
|
COALESCE(m.is_available, TRUE) AS is_available,
|
|
m.config,
|
|
EXTRACT(EPOCH FROM m.created_at)::bigint AS created_at_unix_ms,
|
|
EXTRACT(EPOCH FROM m.updated_at)::bigint AS updated_at_unix_secs,
|
|
gm.name AS global_model_name,
|
|
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
|
|
WHERE m.provider_id = $1
|
|
AND m.is_active = TRUE
|
|
AND gm.is_active = TRUE
|
|
ORDER BY gm.name ASC, m.created_at DESC, m.id ASC
|
|
"#,
|
|
)
|
|
.bind(provider_id)
|
|
.fetch(&self.pool),
|
|
map_admin_provider_model_row,
|
|
)
|
|
.await
|
|
}
|
|
|
|
pub async fn get_admin_global_model_by_id(
|
|
&self,
|
|
global_model_id: &str,
|
|
) -> Result<Option<StoredAdminGlobalModel>, DataLayerError> {
|
|
let row = sqlx::query(
|
|
r#"
|
|
SELECT
|
|
gm.id,
|
|
gm.name,
|
|
COALESCE(NULLIF(gm.display_name, ''), gm.name) AS display_name,
|
|
gm.is_active,
|
|
CAST(gm.default_price_per_request AS DOUBLE PRECISION) AS default_price_per_request,
|
|
gm.default_tiered_pricing,
|
|
gm.supported_capabilities,
|
|
gm.config,
|
|
COALESCE(gm_stats.provider_count, 0) AS provider_count,
|
|
COALESCE(gm_stats.active_provider_count, 0) AS active_provider_count,
|
|
COALESCE(gm.usage_count, 0)::bigint AS usage_count,
|
|
EXTRACT(EPOCH FROM gm.created_at)::bigint AS created_at_unix_ms,
|
|
EXTRACT(EPOCH FROM gm.updated_at)::bigint AS updated_at_unix_secs
|
|
FROM global_models gm
|
|
LEFT JOIN (
|
|
SELECT
|
|
m.global_model_id,
|
|
COUNT(DISTINCT m.provider_id)::bigint AS provider_count,
|
|
COUNT(
|
|
DISTINCT CASE
|
|
WHEN m.is_active = TRUE AND COALESCE(m.is_available, TRUE) = TRUE AND p.is_active = TRUE THEN m.provider_id
|
|
ELSE NULL
|
|
END
|
|
)::bigint AS active_provider_count
|
|
FROM models m
|
|
JOIN providers p ON p.id = m.provider_id
|
|
GROUP BY m.global_model_id
|
|
) gm_stats ON gm_stats.global_model_id = gm.id
|
|
WHERE gm.id = $1
|
|
LIMIT 1
|
|
"#,
|
|
)
|
|
.bind(global_model_id)
|
|
.fetch_optional(&self.pool)
|
|
.await
|
|
.map_postgres_err()?;
|
|
|
|
row.as_ref().map(map_admin_global_model_row).transpose()
|
|
}
|
|
|
|
pub async fn get_admin_global_model_by_name(
|
|
&self,
|
|
model_name: &str,
|
|
) -> Result<Option<StoredAdminGlobalModel>, DataLayerError> {
|
|
let row = sqlx::query(
|
|
r#"
|
|
SELECT
|
|
gm.id,
|
|
gm.name,
|
|
COALESCE(NULLIF(gm.display_name, ''), gm.name) AS display_name,
|
|
gm.is_active,
|
|
CAST(gm.default_price_per_request AS DOUBLE PRECISION) AS default_price_per_request,
|
|
gm.default_tiered_pricing,
|
|
gm.supported_capabilities,
|
|
gm.config,
|
|
COALESCE(gm_stats.provider_count, 0) AS provider_count,
|
|
COALESCE(gm_stats.active_provider_count, 0) AS active_provider_count,
|
|
COALESCE(gm.usage_count, 0)::bigint AS usage_count,
|
|
EXTRACT(EPOCH FROM gm.created_at)::bigint AS created_at_unix_ms,
|
|
EXTRACT(EPOCH FROM gm.updated_at)::bigint AS updated_at_unix_secs
|
|
FROM global_models gm
|
|
LEFT JOIN (
|
|
SELECT
|
|
m.global_model_id,
|
|
COUNT(DISTINCT m.provider_id)::bigint AS provider_count,
|
|
COUNT(
|
|
DISTINCT CASE
|
|
WHEN m.is_active = TRUE AND COALESCE(m.is_available, TRUE) = TRUE AND p.is_active = TRUE THEN m.provider_id
|
|
ELSE NULL
|
|
END
|
|
)::bigint AS active_provider_count
|
|
FROM models m
|
|
JOIN providers p ON p.id = m.provider_id
|
|
GROUP BY m.global_model_id
|
|
) gm_stats ON gm_stats.global_model_id = gm.id
|
|
WHERE gm.name = $1
|
|
LIMIT 1
|
|
"#,
|
|
)
|
|
.bind(model_name)
|
|
.fetch_optional(&self.pool)
|
|
.await
|
|
.map_postgres_err()?;
|
|
|
|
row.as_ref().map(map_admin_global_model_row).transpose()
|
|
}
|
|
|
|
pub async fn list_admin_provider_models_by_global_model_id(
|
|
&self,
|
|
global_model_id: &str,
|
|
) -> Result<Vec<StoredAdminProviderModel>, DataLayerError> {
|
|
collect_query_rows(
|
|
sqlx::query(
|
|
r#"
|
|
SELECT
|
|
m.id,
|
|
m.provider_id,
|
|
m.global_model_id,
|
|
m.provider_model_name,
|
|
m.provider_model_mappings,
|
|
CAST(m.price_per_request AS DOUBLE PRECISION) AS price_per_request,
|
|
m.tiered_pricing,
|
|
m.supports_vision,
|
|
m.supports_function_calling,
|
|
m.supports_streaming,
|
|
m.supports_extended_thinking,
|
|
m.supports_image_generation,
|
|
m.is_active,
|
|
COALESCE(m.is_available, TRUE) AS is_available,
|
|
m.config,
|
|
EXTRACT(EPOCH FROM m.created_at)::bigint AS created_at_unix_ms,
|
|
EXTRACT(EPOCH FROM m.updated_at)::bigint AS updated_at_unix_secs,
|
|
gm.name AS global_model_name,
|
|
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
|
|
WHERE m.global_model_id = $1
|
|
ORDER BY m.created_at DESC, m.id ASC
|
|
"#,
|
|
)
|
|
.bind(global_model_id)
|
|
.fetch(&self.pool),
|
|
map_admin_provider_model_row,
|
|
)
|
|
.await
|
|
}
|
|
|
|
pub async fn create_admin_provider_model(
|
|
&self,
|
|
record: &UpsertAdminProviderModelRecord,
|
|
) -> Result<Option<StoredAdminProviderModel>, DataLayerError> {
|
|
let inserted = sqlx::query(
|
|
r#"
|
|
INSERT INTO models (
|
|
id,
|
|
provider_id,
|
|
global_model_id,
|
|
provider_model_name,
|
|
provider_model_mappings,
|
|
price_per_request,
|
|
tiered_pricing,
|
|
supports_vision,
|
|
supports_function_calling,
|
|
supports_streaming,
|
|
supports_extended_thinking,
|
|
supports_image_generation,
|
|
is_active,
|
|
is_available,
|
|
config,
|
|
created_at,
|
|
updated_at
|
|
)
|
|
VALUES (
|
|
$1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, NOW(), NOW()
|
|
)
|
|
RETURNING id
|
|
"#,
|
|
)
|
|
.bind(&record.id)
|
|
.bind(&record.provider_id)
|
|
.bind(&record.global_model_id)
|
|
.bind(&record.provider_model_name)
|
|
.bind(record.provider_model_mappings.clone())
|
|
.bind(record.price_per_request)
|
|
.bind(record.tiered_pricing.clone())
|
|
.bind(record.supports_vision)
|
|
.bind(record.supports_function_calling)
|
|
.bind(record.supports_streaming)
|
|
.bind(record.supports_extended_thinking)
|
|
.bind(record.supports_image_generation)
|
|
.bind(record.is_active)
|
|
.bind(record.is_available)
|
|
.bind(record.config.clone())
|
|
.fetch_optional(&self.pool)
|
|
.await
|
|
.map_postgres_err()?;
|
|
|
|
if inserted.is_none() {
|
|
return Ok(None);
|
|
}
|
|
|
|
self.get_admin_provider_model(&record.provider_id, &record.id)
|
|
.await
|
|
}
|
|
|
|
pub async fn update_admin_provider_model(
|
|
&self,
|
|
record: &UpsertAdminProviderModelRecord,
|
|
) -> Result<Option<StoredAdminProviderModel>, DataLayerError> {
|
|
let updated = sqlx::query(
|
|
r#"
|
|
UPDATE models
|
|
SET
|
|
global_model_id = $3,
|
|
provider_model_name = $4,
|
|
provider_model_mappings = $5,
|
|
price_per_request = $6,
|
|
tiered_pricing = $7,
|
|
supports_vision = $8,
|
|
supports_function_calling = $9,
|
|
supports_streaming = $10,
|
|
supports_extended_thinking = $11,
|
|
supports_image_generation = $12,
|
|
is_active = $13,
|
|
is_available = $14,
|
|
config = $15,
|
|
updated_at = NOW()
|
|
WHERE id = $1
|
|
AND provider_id = $2
|
|
RETURNING id
|
|
"#,
|
|
)
|
|
.bind(&record.id)
|
|
.bind(&record.provider_id)
|
|
.bind(&record.global_model_id)
|
|
.bind(&record.provider_model_name)
|
|
.bind(record.provider_model_mappings.clone())
|
|
.bind(record.price_per_request)
|
|
.bind(record.tiered_pricing.clone())
|
|
.bind(record.supports_vision)
|
|
.bind(record.supports_function_calling)
|
|
.bind(record.supports_streaming)
|
|
.bind(record.supports_extended_thinking)
|
|
.bind(record.supports_image_generation)
|
|
.bind(record.is_active)
|
|
.bind(record.is_available)
|
|
.bind(record.config.clone())
|
|
.fetch_optional(&self.pool)
|
|
.await
|
|
.map_postgres_err()?;
|
|
|
|
if updated.is_none() {
|
|
return Ok(None);
|
|
}
|
|
|
|
self.get_admin_provider_model(&record.provider_id, &record.id)
|
|
.await
|
|
}
|
|
|
|
pub async fn delete_admin_provider_model(
|
|
&self,
|
|
provider_id: &str,
|
|
model_id: &str,
|
|
) -> Result<bool, DataLayerError> {
|
|
let deleted = sqlx::query(
|
|
r#"
|
|
DELETE FROM models
|
|
WHERE provider_id = $1
|
|
AND id = $2
|
|
RETURNING id
|
|
"#,
|
|
)
|
|
.bind(provider_id)
|
|
.bind(model_id)
|
|
.fetch_optional(&self.pool)
|
|
.await
|
|
.map_postgres_err()?;
|
|
|
|
Ok(deleted.is_some())
|
|
}
|
|
|
|
pub async fn create_admin_global_model(
|
|
&self,
|
|
record: &CreateAdminGlobalModelRecord,
|
|
) -> Result<Option<StoredAdminGlobalModel>, DataLayerError> {
|
|
let usage_count =
|
|
optional_admin_global_model_usage_count_i64(record.usage_count)?.unwrap_or_default();
|
|
let inserted = sqlx::query(
|
|
r#"
|
|
INSERT INTO global_models (
|
|
id,
|
|
name,
|
|
display_name,
|
|
is_active,
|
|
default_price_per_request,
|
|
default_tiered_pricing,
|
|
supported_capabilities,
|
|
usage_count,
|
|
config,
|
|
created_at,
|
|
updated_at
|
|
)
|
|
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, NOW(), NOW())
|
|
RETURNING id
|
|
"#,
|
|
)
|
|
.bind(&record.id)
|
|
.bind(&record.name)
|
|
.bind(&record.display_name)
|
|
.bind(record.is_active)
|
|
.bind(record.default_price_per_request)
|
|
.bind(record.default_tiered_pricing.clone())
|
|
.bind(record.supported_capabilities.clone())
|
|
.bind(usage_count)
|
|
.bind(record.config.clone())
|
|
.fetch_optional(&self.pool)
|
|
.await
|
|
.map_postgres_err()?;
|
|
|
|
if inserted.is_none() {
|
|
return Ok(None);
|
|
}
|
|
|
|
self.get_admin_global_model_by_id(&record.id).await
|
|
}
|
|
|
|
pub async fn update_admin_global_model(
|
|
&self,
|
|
record: &UpdateAdminGlobalModelRecord,
|
|
) -> Result<Option<StoredAdminGlobalModel>, DataLayerError> {
|
|
let usage_count = optional_admin_global_model_usage_count_i64(record.usage_count)?;
|
|
let updated = sqlx::query(
|
|
r#"
|
|
UPDATE global_models
|
|
SET
|
|
display_name = $2,
|
|
is_active = $3,
|
|
default_price_per_request = $4,
|
|
default_tiered_pricing = $5,
|
|
supported_capabilities = $6,
|
|
config = $7,
|
|
usage_count = COALESCE($8, usage_count),
|
|
updated_at = NOW()
|
|
WHERE id = $1
|
|
RETURNING id
|
|
"#,
|
|
)
|
|
.bind(&record.id)
|
|
.bind(&record.display_name)
|
|
.bind(record.is_active)
|
|
.bind(record.default_price_per_request)
|
|
.bind(record.default_tiered_pricing.clone())
|
|
.bind(record.supported_capabilities.clone())
|
|
.bind(record.config.clone())
|
|
.bind(usage_count)
|
|
.fetch_optional(&self.pool)
|
|
.await
|
|
.map_postgres_err()?;
|
|
|
|
if updated.is_none() {
|
|
return Ok(None);
|
|
}
|
|
|
|
self.get_admin_global_model_by_id(&record.id).await
|
|
}
|
|
|
|
pub async fn delete_admin_global_model(
|
|
&self,
|
|
global_model_id: &str,
|
|
) -> Result<bool, DataLayerError> {
|
|
let mut tx = self.pool.begin().await.map_postgres_err()?;
|
|
|
|
sqlx::query(
|
|
r#"
|
|
DELETE FROM models
|
|
WHERE global_model_id = $1
|
|
"#,
|
|
)
|
|
.bind(global_model_id)
|
|
.execute(&mut *tx)
|
|
.await
|
|
.map_postgres_err()?;
|
|
|
|
let deleted = sqlx::query(
|
|
r#"
|
|
DELETE FROM global_models
|
|
WHERE id = $1
|
|
RETURNING id
|
|
"#,
|
|
)
|
|
.bind(global_model_id)
|
|
.fetch_optional(&mut *tx)
|
|
.await
|
|
.map_postgres_err()?;
|
|
|
|
tx.commit().await.map_postgres_err()?;
|
|
|
|
Ok(deleted.is_some())
|
|
}
|
|
}
|
|
|
|
#[async_trait]
|
|
impl GlobalModelReadRepository for SqlxGlobalModelReadRepository {
|
|
async fn list_public_models(
|
|
&self,
|
|
query: &PublicGlobalModelQuery,
|
|
) -> Result<StoredPublicGlobalModelPage, DataLayerError> {
|
|
Self::list_public_models(self, query).await
|
|
}
|
|
|
|
async fn get_public_model_by_name(
|
|
&self,
|
|
model_name: &str,
|
|
) -> Result<Option<StoredPublicGlobalModel>, DataLayerError> {
|
|
let row = sqlx::query(
|
|
r#"
|
|
SELECT
|
|
id,
|
|
name,
|
|
display_name,
|
|
is_active,
|
|
CAST(default_price_per_request AS DOUBLE PRECISION) AS default_price_per_request,
|
|
default_tiered_pricing,
|
|
supported_capabilities,
|
|
config
|
|
FROM global_models
|
|
WHERE name = $1 AND is_active = TRUE
|
|
LIMIT 1
|
|
"#,
|
|
)
|
|
.bind(model_name)
|
|
.fetch_optional(&self.pool)
|
|
.await
|
|
.map_postgres_err()?;
|
|
|
|
row.as_ref().map(map_row).transpose()
|
|
}
|
|
|
|
async fn list_public_catalog_models(
|
|
&self,
|
|
query: &PublicCatalogModelListQuery,
|
|
) -> Result<Vec<StoredPublicCatalogModel>, DataLayerError> {
|
|
let mut builder = QueryBuilder::<Postgres>::new(LIST_PUBLIC_CATALOG_MODELS_PREFIX);
|
|
apply_public_catalog_model_filters(&mut builder, query.provider_id.as_deref(), None);
|
|
builder
|
|
.push(" ORDER BY p.provider_priority ASC, p.name ASC, COALESCE(gm.name, m.provider_model_name) ASC, m.id ASC OFFSET ")
|
|
.push_bind(query.offset as i64)
|
|
.push(" LIMIT ")
|
|
.push_bind(query.limit as i64);
|
|
let query = builder.build();
|
|
collect_query_rows(query.fetch(&self.pool), map_public_catalog_model_row).await
|
|
}
|
|
|
|
async fn search_public_catalog_models(
|
|
&self,
|
|
query: &PublicCatalogModelSearchQuery,
|
|
) -> Result<Vec<StoredPublicCatalogModel>, DataLayerError> {
|
|
let mut builder = QueryBuilder::<Postgres>::new(LIST_PUBLIC_CATALOG_MODELS_PREFIX);
|
|
apply_public_catalog_model_filters(
|
|
&mut builder,
|
|
query.provider_id.as_deref(),
|
|
Some(query.search.as_str()),
|
|
);
|
|
builder
|
|
.push(" ORDER BY p.provider_priority ASC, p.name ASC, COALESCE(gm.name, m.provider_model_name) ASC, m.id ASC LIMIT ")
|
|
.push_bind(query.limit as i64);
|
|
let query = builder.build();
|
|
collect_query_rows(query.fetch(&self.pool), map_public_catalog_model_row).await
|
|
}
|
|
|
|
async fn list_admin_global_models(
|
|
&self,
|
|
query: &AdminGlobalModelListQuery,
|
|
) -> Result<StoredAdminGlobalModelPage, DataLayerError> {
|
|
Self::list_admin_global_models(self, query).await
|
|
}
|
|
|
|
async fn list_admin_provider_models(
|
|
&self,
|
|
query: &AdminProviderModelListQuery,
|
|
) -> Result<Vec<StoredAdminProviderModel>, DataLayerError> {
|
|
Self::list_admin_provider_models(self, query).await
|
|
}
|
|
|
|
async fn get_admin_provider_model(
|
|
&self,
|
|
provider_id: &str,
|
|
model_id: &str,
|
|
) -> Result<Option<StoredAdminProviderModel>, DataLayerError> {
|
|
Self::get_admin_provider_model(self, provider_id, model_id).await
|
|
}
|
|
|
|
async fn list_admin_provider_available_source_models(
|
|
&self,
|
|
provider_id: &str,
|
|
) -> Result<Vec<StoredAdminProviderModel>, DataLayerError> {
|
|
Self::list_admin_provider_available_source_models(self, provider_id).await
|
|
}
|
|
|
|
async fn get_admin_global_model_by_id(
|
|
&self,
|
|
global_model_id: &str,
|
|
) -> Result<Option<StoredAdminGlobalModel>, DataLayerError> {
|
|
Self::get_admin_global_model_by_id(self, global_model_id).await
|
|
}
|
|
|
|
async fn get_admin_global_model_by_name(
|
|
&self,
|
|
model_name: &str,
|
|
) -> Result<Option<StoredAdminGlobalModel>, DataLayerError> {
|
|
Self::get_admin_global_model_by_name(self, model_name).await
|
|
}
|
|
|
|
async fn list_admin_provider_models_by_global_model_id(
|
|
&self,
|
|
global_model_id: &str,
|
|
) -> Result<Vec<StoredAdminProviderModel>, DataLayerError> {
|
|
Self::list_admin_provider_models_by_global_model_id(self, global_model_id).await
|
|
}
|
|
|
|
async fn list_provider_model_stats(
|
|
&self,
|
|
provider_ids: &[String],
|
|
) -> Result<Vec<StoredProviderModelStats>, DataLayerError> {
|
|
Self::list_provider_model_stats(self, provider_ids).await
|
|
}
|
|
|
|
async fn list_active_global_model_ids_by_provider_ids(
|
|
&self,
|
|
provider_ids: &[String],
|
|
) -> Result<Vec<StoredProviderActiveGlobalModel>, DataLayerError> {
|
|
Self::list_active_global_model_ids_by_provider_ids(self, provider_ids).await
|
|
}
|
|
}
|
|
|
|
#[async_trait]
|
|
impl GlobalModelWriteRepository for SqlxGlobalModelReadRepository {
|
|
async fn create_admin_provider_model(
|
|
&self,
|
|
record: &UpsertAdminProviderModelRecord,
|
|
) -> Result<Option<StoredAdminProviderModel>, DataLayerError> {
|
|
Self::create_admin_provider_model(self, record).await
|
|
}
|
|
|
|
async fn update_admin_provider_model(
|
|
&self,
|
|
record: &UpsertAdminProviderModelRecord,
|
|
) -> Result<Option<StoredAdminProviderModel>, DataLayerError> {
|
|
Self::update_admin_provider_model(self, record).await
|
|
}
|
|
|
|
async fn delete_admin_provider_model(
|
|
&self,
|
|
provider_id: &str,
|
|
model_id: &str,
|
|
) -> Result<bool, DataLayerError> {
|
|
Self::delete_admin_provider_model(self, provider_id, model_id).await
|
|
}
|
|
|
|
async fn create_admin_global_model(
|
|
&self,
|
|
record: &CreateAdminGlobalModelRecord,
|
|
) -> Result<Option<StoredAdminGlobalModel>, DataLayerError> {
|
|
Self::create_admin_global_model(self, record).await
|
|
}
|
|
|
|
async fn update_admin_global_model(
|
|
&self,
|
|
record: &UpdateAdminGlobalModelRecord,
|
|
) -> Result<Option<StoredAdminGlobalModel>, DataLayerError> {
|
|
Self::update_admin_global_model(self, record).await
|
|
}
|
|
|
|
async fn delete_admin_global_model(
|
|
&self,
|
|
global_model_id: &str,
|
|
) -> Result<bool, DataLayerError> {
|
|
Self::delete_admin_global_model(self, global_model_id).await
|
|
}
|
|
}
|
|
|
|
fn apply_public_model_filters(
|
|
builder: &mut QueryBuilder<'_, Postgres>,
|
|
query: &PublicGlobalModelQuery,
|
|
) {
|
|
builder.push(" WHERE ");
|
|
match query.is_active {
|
|
Some(is_active) => {
|
|
builder.push("is_active = ").push_bind(is_active);
|
|
}
|
|
None => {
|
|
builder.push("is_active = TRUE");
|
|
}
|
|
}
|
|
|
|
if let Some(search) = query
|
|
.search
|
|
.as_deref()
|
|
.map(str::trim)
|
|
.filter(|value| !value.is_empty())
|
|
{
|
|
let pattern = format!("%{search}%");
|
|
builder
|
|
.push(" AND (name ILIKE ")
|
|
.push_bind(pattern.clone())
|
|
.push(" OR display_name ILIKE ")
|
|
.push_bind(pattern)
|
|
.push(")");
|
|
}
|
|
}
|
|
|
|
fn apply_admin_global_model_filters(
|
|
builder: &mut QueryBuilder<'_, Postgres>,
|
|
query: &AdminGlobalModelListQuery,
|
|
) {
|
|
builder.push(" WHERE 1=1");
|
|
if let Some(is_active) = query.is_active {
|
|
builder.push(" AND gm.is_active = ").push_bind(is_active);
|
|
}
|
|
if let Some(search) = query
|
|
.search
|
|
.as_deref()
|
|
.map(str::trim)
|
|
.filter(|value| !value.is_empty())
|
|
{
|
|
let pattern = format!("%{search}%");
|
|
builder
|
|
.push(" AND (gm.name ILIKE ")
|
|
.push_bind(pattern.clone())
|
|
.push(" OR gm.display_name ILIKE ")
|
|
.push_bind(pattern)
|
|
.push(")");
|
|
}
|
|
}
|
|
|
|
fn map_row(row: &PgRow) -> Result<StoredPublicGlobalModel, DataLayerError> {
|
|
let supported_capabilities: Option<Value> =
|
|
row.try_get("supported_capabilities").map_postgres_err()?;
|
|
StoredPublicGlobalModel::new(
|
|
row.try_get("id").map_postgres_err()?,
|
|
row.try_get("name").map_postgres_err()?,
|
|
row.try_get("display_name").map_postgres_err()?,
|
|
row.try_get("is_active").map_postgres_err()?,
|
|
row.try_get("default_price_per_request")
|
|
.map_postgres_err()?,
|
|
row.try_get("default_tiered_pricing").map_postgres_err()?,
|
|
supported_capabilities,
|
|
row.try_get("config").map_postgres_err()?,
|
|
0,
|
|
)
|
|
}
|
|
|
|
async fn collect_query_rows<T, S>(
|
|
mut rows: S,
|
|
map_row: fn(&PgRow) -> Result<T, DataLayerError>,
|
|
) -> Result<Vec<T>, DataLayerError>
|
|
where
|
|
S: TryStream<Ok = PgRow, Error = sqlx::Error> + Unpin,
|
|
{
|
|
let mut items = Vec::new();
|
|
while let Some(row) = rows.try_next().await.map_postgres_err()? {
|
|
items.push(map_row(&row)?);
|
|
}
|
|
Ok(items)
|
|
}
|
|
|
|
fn apply_public_catalog_model_filters(
|
|
builder: &mut QueryBuilder<'_, Postgres>,
|
|
provider_id: Option<&str>,
|
|
search: Option<&str>,
|
|
) {
|
|
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
|
|
.push(" AND m.provider_id = ")
|
|
.push_bind(provider_id.to_string());
|
|
}
|
|
|
|
if let Some(search) = search.map(str::trim).filter(|value| !value.is_empty()) {
|
|
let pattern = format!("%{search}%");
|
|
builder
|
|
.push(" AND (m.provider_model_name ILIKE ")
|
|
.push_bind(pattern.clone())
|
|
.push(" OR gm.name ILIKE ")
|
|
.push_bind(pattern.clone())
|
|
.push(" OR gm.display_name ILIKE ")
|
|
.push_bind(pattern)
|
|
.push(")");
|
|
}
|
|
}
|
|
|
|
fn map_public_catalog_model_row(row: &PgRow) -> Result<StoredPublicCatalogModel, DataLayerError> {
|
|
StoredPublicCatalogModel::new(
|
|
row.try_get("id").map_postgres_err()?,
|
|
row.try_get("provider_id").map_postgres_err()?,
|
|
row.try_get("provider_name").map_postgres_err()?,
|
|
row.try_get("provider_model_name").map_postgres_err()?,
|
|
row.try_get("name").map_postgres_err()?,
|
|
row.try_get("display_name").map_postgres_err()?,
|
|
row.try_get("description").map_postgres_err()?,
|
|
row.try_get("icon_url").map_postgres_err()?,
|
|
row.try_get("input_price_per_1m").map_postgres_err()?,
|
|
row.try_get("output_price_per_1m").map_postgres_err()?,
|
|
row.try_get("cache_creation_price_per_1m")
|
|
.map_postgres_err()?,
|
|
row.try_get("cache_read_price_per_1m").map_postgres_err()?,
|
|
row.try_get("supports_vision").map_postgres_err()?,
|
|
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()?,
|
|
)
|
|
}
|
|
|
|
fn map_admin_provider_model_row(row: &PgRow) -> Result<StoredAdminProviderModel, DataLayerError> {
|
|
let created_at_unix_ms = row
|
|
.try_get::<Option<i64>, _>("created_at_unix_ms")
|
|
.map_postgres_err()?
|
|
.map(|value| value.max(0) as u64);
|
|
let updated_at_unix_secs = row
|
|
.try_get::<Option<i64>, _>("updated_at_unix_secs")
|
|
.map_postgres_err()?
|
|
.map(|value| value.max(0) as u64);
|
|
StoredAdminProviderModel::new(
|
|
row.try_get("id").map_postgres_err()?,
|
|
row.try_get("provider_id").map_postgres_err()?,
|
|
row.try_get("global_model_id").map_postgres_err()?,
|
|
row.try_get("provider_model_name").map_postgres_err()?,
|
|
row.try_get("provider_model_mappings").map_postgres_err()?,
|
|
row.try_get("price_per_request").map_postgres_err()?,
|
|
row.try_get("tiered_pricing").map_postgres_err()?,
|
|
row.try_get("supports_vision").map_postgres_err()?,
|
|
row.try_get("supports_function_calling")
|
|
.map_postgres_err()?,
|
|
row.try_get("supports_streaming").map_postgres_err()?,
|
|
row.try_get("supports_extended_thinking")
|
|
.map_postgres_err()?,
|
|
row.try_get("supports_image_generation")
|
|
.map_postgres_err()?,
|
|
row.try_get("is_active").map_postgres_err()?,
|
|
row.try_get("is_available").map_postgres_err()?,
|
|
row.try_get("config").map_postgres_err()?,
|
|
created_at_unix_ms,
|
|
updated_at_unix_secs,
|
|
row.try_get("global_model_name").map_postgres_err()?,
|
|
row.try_get("global_model_display_name")
|
|
.map_postgres_err()?,
|
|
row.try_get("global_model_default_price_per_request")
|
|
.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()?,
|
|
)
|
|
}
|
|
|
|
fn map_admin_global_model_row(row: &PgRow) -> Result<StoredAdminGlobalModel, DataLayerError> {
|
|
let created_at_unix_ms = row
|
|
.try_get::<Option<i64>, _>("created_at_unix_ms")
|
|
.map_postgres_err()?
|
|
.map(|value| value.max(0) as u64);
|
|
let updated_at_unix_secs = row
|
|
.try_get::<Option<i64>, _>("updated_at_unix_secs")
|
|
.map_postgres_err()?
|
|
.map(|value| value.max(0) as u64);
|
|
let provider_count = row
|
|
.try_get::<i64, _>("provider_count")
|
|
.map_postgres_err()?
|
|
.max(0) as u64;
|
|
let active_provider_count = row
|
|
.try_get::<i64, _>("active_provider_count")
|
|
.map_postgres_err()?
|
|
.max(0) as u64;
|
|
let usage_count = row
|
|
.try_get::<i64, _>("usage_count")
|
|
.map_postgres_err()?
|
|
.max(0) as u64;
|
|
StoredAdminGlobalModel::new(
|
|
row.try_get("id").map_postgres_err()?,
|
|
row.try_get("name").map_postgres_err()?,
|
|
row.try_get("display_name").map_postgres_err()?,
|
|
row.try_get("is_active").map_postgres_err()?,
|
|
row.try_get("default_price_per_request")
|
|
.map_postgres_err()?,
|
|
row.try_get("default_tiered_pricing").map_postgres_err()?,
|
|
row.try_get("supported_capabilities").map_postgres_err()?,
|
|
row.try_get("config").map_postgres_err()?,
|
|
provider_count,
|
|
active_provider_count,
|
|
usage_count,
|
|
created_at_unix_ms,
|
|
updated_at_unix_secs,
|
|
)
|
|
}
|
|
|
|
fn build_provider_id_list_query<'a>(
|
|
prefix: &'static str,
|
|
provider_ids: &'a [String],
|
|
suffix: &'static str,
|
|
) -> QueryBuilder<'a, Postgres> {
|
|
let mut builder = QueryBuilder::<Postgres>::new(prefix);
|
|
let mut separated = builder.separated(", ");
|
|
for provider_id in provider_ids {
|
|
separated.push_bind(provider_id);
|
|
}
|
|
separated.push_unseparated(suffix);
|
|
builder
|
|
}
|
|
|
|
fn map_provider_model_stats_row(row: &PgRow) -> Result<StoredProviderModelStats, DataLayerError> {
|
|
StoredProviderModelStats::new(
|
|
row.try_get("provider_id").map_postgres_err()?,
|
|
row.try_get("total_models").map_postgres_err()?,
|
|
row.try_get("active_models").map_postgres_err()?,
|
|
)
|
|
}
|
|
|
|
fn map_provider_active_global_model_row(
|
|
row: &PgRow,
|
|
) -> Result<StoredProviderActiveGlobalModel, DataLayerError> {
|
|
StoredProviderActiveGlobalModel::new(
|
|
row.try_get("provider_id").map_postgres_err()?,
|
|
row.try_get("global_model_id").map_postgres_err()?,
|
|
)
|
|
}
|
|
|
|
fn optional_admin_global_model_usage_count_i64(
|
|
value: Option<u64>,
|
|
) -> Result<Option<i64>, DataLayerError> {
|
|
value
|
|
.map(|value| {
|
|
i64::try_from(value).map_err(|_| {
|
|
DataLayerError::InvalidInput(
|
|
"global_models.usage_count exceeds i64 range".to_string(),
|
|
)
|
|
})
|
|
})
|
|
.transpose()
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::{
|
|
SqlxGlobalModelReadRepository, LIST_ADMIN_GLOBAL_MODELS_PREFIX,
|
|
LIST_ADMIN_PROVIDER_MODELS_PREFIX,
|
|
};
|
|
use crate::{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!(
|
|
"global_models.rs"
|
|
));
|
|
let supported_capabilities_projection = format!(
|
|
"{} AS {}",
|
|
"gm.supported_capabilities", "global_model_supported_capabilities"
|
|
);
|
|
assert_eq!(
|
|
include_str!("global_models.rs")
|
|
.matches(&supported_capabilities_projection)
|
|
.count(),
|
|
4
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn admin_global_model_sql_projects_usage_count_from_billing_facts() {
|
|
let source = include_str!("global_models.rs");
|
|
|
|
assert!(
|
|
!LIST_ADMIN_GLOBAL_MODELS_PREFIX.contains("usage_billing_facts"),
|
|
"admin global model list should read the maintained usage_count field"
|
|
);
|
|
assert!(
|
|
source.contains("WHERE usage.model = gm.name"),
|
|
"admin global model detail lookups should count usage facts for the selected model"
|
|
);
|
|
assert!(
|
|
source.contains("AND usage.status NOT IN ('pending', 'streaming')"),
|
|
"admin global model detail usage_count should use the maintained read-model status scope"
|
|
);
|
|
assert!(
|
|
source.contains("COALESCE(usage_stats.usage_count, gm.usage_count, 0)::bigint AS usage_count"),
|
|
"admin global model usage_count should prefer actual usage facts with stored-count fallback"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn admin_global_model_list_sql_counts_usage_without_full_fact_aggregate() {
|
|
assert!(
|
|
LIST_ADMIN_GLOBAL_MODELS_PREFIX.contains("COALESCE(gm.usage_count, 0)::bigint AS usage_count"),
|
|
"admin global model list should read maintained usage_count without per-request fact scans"
|
|
);
|
|
assert!(
|
|
!LIST_ADMIN_GLOBAL_MODELS_PREFIX.contains("GROUP BY usage.model"),
|
|
"admin global model list must not aggregate the full usage fact table before pagination"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn global_model_usage_count_read_model_has_backfill_and_delta_maintenance() {
|
|
let backfill_sql = include_str!(
|
|
"../../../runtime/backfills/postgres/20260505120000_rebuild_global_model_usage_count.sql"
|
|
);
|
|
let delta_sql = include_str!("usage/queries/apply_global_model_usage_delta_sql.sql");
|
|
|
|
assert!(
|
|
backfill_sql.contains("UPDATE global_models AS gm"),
|
|
"global model usage_count backfill should refresh the read model"
|
|
);
|
|
assert!(
|
|
backfill_sql.contains("FROM usage_billing_facts AS usage"),
|
|
"global model usage_count backfill should rebuild from canonical usage facts"
|
|
);
|
|
assert!(
|
|
delta_sql.contains("UPDATE global_models"),
|
|
"usage writes should maintain the global model usage_count read model"
|
|
);
|
|
assert!(
|
|
delta_sql.contains("WHERE name = $1"),
|
|
"global model usage_count delta should target models by canonical model name"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn repository_constructs_from_lazy_pool() {
|
|
let factory = PostgresPoolFactory::new(PostgresPoolConfig {
|
|
database_url: "postgres://localhost/aether".to_string(),
|
|
min_connections: 1,
|
|
max_connections: 4,
|
|
acquire_timeout_ms: 1_000,
|
|
idle_timeout_ms: 5_000,
|
|
max_lifetime_ms: 30_000,
|
|
statement_cache_capacity: 64,
|
|
require_ssl: false,
|
|
})
|
|
.expect("factory should build");
|
|
|
|
let pool = factory.connect_lazy().expect("pool should build");
|
|
let _repository = SqlxGlobalModelReadRepository::new(pool);
|
|
}
|
|
}
|