fix(admin): align global model responses with repository counts

This commit is contained in:
fawney19
2026-04-11 02:55:02 +08:00
parent a54ac76688
commit 34d295c1e0
6 changed files with 91 additions and 132 deletions

View File

@@ -6,9 +6,7 @@ use crate::handlers::admin::request::AdminAppState;
use aether_data_contracts::repository::global_models::{ use aether_data_contracts::repository::global_models::{
StoredAdminGlobalModel, StoredAdminProviderModel, StoredAdminGlobalModel, StoredAdminProviderModel,
}; };
use futures_util::stream::{self, StreamExt};
use serde_json::json; use serde_json::json;
use std::collections::{BTreeMap, BTreeSet};
use std::time::{SystemTime, UNIX_EPOCH}; use std::time::{SystemTime, UNIX_EPOCH};
pub(crate) async fn resolve_admin_global_model_by_id_or_err( pub(crate) async fn resolve_admin_global_model_by_id_or_err(
@@ -30,24 +28,6 @@ pub(super) fn admin_global_models_now_unix_secs() -> u64 {
.unwrap_or(0) .unwrap_or(0)
} }
pub(super) fn admin_global_model_provider_counts(
provider_models: &[StoredAdminProviderModel],
) -> (usize, usize, usize) {
let total_models = provider_models.len();
let total_providers = provider_models
.iter()
.map(|model| model.provider_id.clone())
.collect::<BTreeSet<_>>()
.len();
let active_provider_count = provider_models
.iter()
.filter(|model| model.is_active && model.is_available)
.map(|model| model.provider_id.clone())
.collect::<BTreeSet<_>>()
.len();
(total_models, total_providers, active_provider_count)
}
pub(super) fn build_admin_global_model_price_range( pub(super) fn build_admin_global_model_price_range(
global_model: &StoredAdminGlobalModel, global_model: &StoredAdminGlobalModel,
provider_models: &[StoredAdminProviderModel], provider_models: &[StoredAdminProviderModel],
@@ -85,27 +65,3 @@ pub(super) fn build_admin_global_model_price_range(
"max_output": output_values.iter().copied().reduce(f64::max), "max_output": output_values.iter().copied().reduce(f64::max),
}) })
} }
#[allow(clippy::redundant_iter_cloned, clippy::redundant_locals)]
pub(super) async fn admin_global_model_provider_models_by_global_model_id(
state: &AdminAppState<'_>,
global_model_ids: &[String],
) -> BTreeMap<String, Vec<StoredAdminProviderModel>> {
let state = *state;
stream::iter(global_model_ids.iter().cloned().map(|global_model_id| {
let state = state;
async move {
let provider_models = state
.list_admin_provider_models_by_global_model_id(&global_model_id)
.await
.ok()
.unwrap_or_default();
(global_model_id, provider_models)
}
}))
.buffer_unordered(32)
.collect::<Vec<_>>()
.await
.into_iter()
.collect()
}

View File

@@ -1,22 +1,16 @@
use super::super::super::shared::json_string_list; use super::super::super::shared::json_string_list;
use super::super::payloads::timestamp_or_now; use super::super::payloads::timestamp_or_now;
use super::helpers::{ use super::helpers::{admin_global_models_now_unix_secs, build_admin_global_model_price_range};
admin_global_model_provider_counts, admin_global_model_provider_models_by_global_model_id,
admin_global_models_now_unix_secs, build_admin_global_model_price_range,
};
use crate::handlers::admin::request::AdminAppState; use crate::handlers::admin::request::AdminAppState;
use aether_data_contracts::repository::global_models::{ use aether_data_contracts::repository::global_models::{
AdminGlobalModelListQuery, StoredAdminGlobalModel, StoredAdminProviderModel, AdminGlobalModelListQuery, StoredAdminGlobalModel,
}; };
use serde_json::json; use serde_json::json;
pub(crate) fn build_admin_global_model_response( pub(crate) fn build_admin_global_model_response(
global_model: &StoredAdminGlobalModel, global_model: &StoredAdminGlobalModel,
provider_models: &[StoredAdminProviderModel],
now_unix_secs: u64, now_unix_secs: u64,
) -> serde_json::Value { ) -> serde_json::Value {
let (_, provider_count, active_provider_count) =
admin_global_model_provider_counts(provider_models);
json!({ json!({
"id": &global_model.id, "id": &global_model.id,
"name": &global_model.name, "name": &global_model.name,
@@ -26,9 +20,9 @@ pub(crate) fn build_admin_global_model_response(
"default_tiered_pricing": global_model.default_tiered_pricing.clone(), "default_tiered_pricing": global_model.default_tiered_pricing.clone(),
"supported_capabilities": json_string_list(global_model.supported_capabilities.as_ref()), "supported_capabilities": json_string_list(global_model.supported_capabilities.as_ref()),
"config": global_model.config.clone(), "config": global_model.config.clone(),
"provider_count": provider_count, "provider_count": global_model.provider_count,
"active_provider_count": active_provider_count, "active_provider_count": global_model.active_provider_count,
"usage_count": 0, "usage_count": global_model.usage_count,
"created_at": timestamp_or_now(global_model.created_at_unix_ms, now_unix_secs), "created_at": timestamp_or_now(global_model.created_at_unix_ms, now_unix_secs),
"updated_at": timestamp_or_now(global_model.updated_at_unix_secs, now_unix_secs), "updated_at": timestamp_or_now(global_model.updated_at_unix_secs, now_unix_secs),
}) })
@@ -60,22 +54,9 @@ pub(crate) async fn build_admin_global_models_payload(
.cmp(&right.name) .cmp(&right.name)
.then_with(|| left.id.cmp(&right.id)) .then_with(|| left.id.cmp(&right.id))
}); });
let global_model_ids = models
.iter()
.map(|model| model.id.clone())
.collect::<Vec<_>>();
let mut provider_models_by_global_model =
admin_global_model_provider_models_by_global_model_id(state, &global_model_ids).await;
let mut payload_models = Vec::with_capacity(models.len()); let mut payload_models = Vec::with_capacity(models.len());
for model in models { for model in models {
let provider_models = provider_models_by_global_model payload_models.push(build_admin_global_model_response(&model, now_unix_secs));
.remove(&model.id)
.unwrap_or_default();
payload_models.push(build_admin_global_model_response(
&model,
&provider_models,
now_unix_secs,
));
} }
Some(json!({ Some(json!({
"models": payload_models, "models": payload_models,
@@ -100,11 +81,11 @@ pub(crate) async fn build_admin_global_model_payload(
.ok() .ok()
.unwrap_or_default(); .unwrap_or_default();
let now_unix_secs = admin_global_models_now_unix_secs(); let now_unix_secs = admin_global_models_now_unix_secs();
let (total_models, total_providers, _) = admin_global_model_provider_counts(&provider_models); let total_models = provider_models.len();
let mut payload = build_admin_global_model_response(&model, &provider_models, now_unix_secs); let mut payload = build_admin_global_model_response(&model, now_unix_secs);
if let Some(object) = payload.as_object_mut() { if let Some(object) = payload.as_object_mut() {
object.insert("total_models".to_string(), json!(total_models)); object.insert("total_models".to_string(), json!(total_models));
object.insert("total_providers".to_string(), json!(total_providers)); object.insert("total_providers".to_string(), json!(model.provider_count));
object.insert( object.insert(
"price_range".to_string(), "price_range".to_string(),
build_admin_global_model_price_range(&model, &provider_models), build_admin_global_model_price_range(&model, &provider_models),

View File

@@ -104,27 +104,20 @@ async fn build_create_global_model_response(
}; };
Ok(match state.create_admin_global_model(&record).await? { Ok(match state.create_admin_global_model(&record).await? {
Some(created) => { Some(created) => attach_admin_audit_response(
let provider_models = state (
.list_admin_provider_models_by_global_model_id(&created.id) http::StatusCode::CREATED,
.await Json(build_admin_global_model_response(
.unwrap_or_default(); &created,
attach_admin_audit_response( current_unix_secs(),
( )),
http::StatusCode::CREATED,
Json(build_admin_global_model_response(
&created,
&provider_models,
current_unix_secs(),
)),
)
.into_response(),
"admin_global_model_created",
"create_global_model",
"global_model",
&created.id,
) )
} .into_response(),
"admin_global_model_created",
"create_global_model",
"global_model",
&created.id,
),
None => ( None => (
http::StatusCode::SERVICE_UNAVAILABLE, http::StatusCode::SERVICE_UNAVAILABLE,
Json(json!({ "detail": ADMIN_GLOBAL_MODELS_DATA_UNAVAILABLE_DETAIL })), Json(json!({ "detail": ADMIN_GLOBAL_MODELS_DATA_UNAVAILABLE_DETAIL })),
@@ -169,24 +162,17 @@ async fn build_update_global_model_response(
}; };
Ok(match state.update_admin_global_model(&record).await? { Ok(match state.update_admin_global_model(&record).await? {
Some(updated) => { Some(updated) => attach_admin_audit_response(
let provider_models = state Json(build_admin_global_model_response(
.list_admin_provider_models_by_global_model_id(&updated.id) &updated,
.await current_unix_secs(),
.unwrap_or_default(); ))
attach_admin_audit_response( .into_response(),
Json(build_admin_global_model_response( "admin_global_model_updated",
&updated, "update_global_model",
&provider_models, "global_model",
current_unix_secs(), &updated.id,
)) ),
.into_response(),
"admin_global_model_updated",
"update_global_model",
"global_model",
&updated.id,
)
}
None => global_model_not_found_response(&existing.id), None => global_model_not_found_response(&existing.id),
}) })
} }

View File

@@ -40,11 +40,27 @@ async fn gateway_handles_admin_global_models_locally_with_trusted_admin_principa
}), }),
); );
let mut gpt41 = sample_admin_global_model("global-gpt-4.1", "gpt-4.1", "GPT 4.1");
gpt41.usage_count = 7;
let mut gpt41_anthropic = sample_admin_provider_model(
"model-anthropic-gpt41",
"provider-anthropic",
"global-gpt-4.1",
"gpt-4.1-anthropic",
);
gpt41_anthropic.is_available = false;
let mut gpt41_google = sample_admin_provider_model(
"model-google-gpt41",
"provider-google",
"global-gpt-4.1",
"gpt-4.1-google",
);
gpt41_google.is_available = false;
let global_model_repository = Arc::new( let global_model_repository = Arc::new(
InMemoryGlobalModelReadRepository::seed(Vec::new()) InMemoryGlobalModelReadRepository::seed(Vec::new())
.with_admin_global_models(vec![ .with_admin_global_models(vec![
sample_admin_global_model("global-gpt-5", "gpt-5", "GPT 5"), sample_admin_global_model("global-gpt-5", "gpt-5", "GPT 5"),
sample_admin_global_model("global-gpt-4.1", "gpt-4.1", "GPT 4.1"), gpt41,
]) ])
.with_admin_provider_models(vec![ .with_admin_provider_models(vec![
sample_admin_provider_model( sample_admin_provider_model(
@@ -59,6 +75,8 @@ async fn gateway_handles_admin_global_models_locally_with_trusted_admin_principa
"global-gpt-4.1", "global-gpt-4.1",
"gpt-4.1-upstream", "gpt-4.1-upstream",
), ),
gpt41_anthropic,
gpt41_google,
]), ]),
); );
@@ -92,9 +110,9 @@ async fn gateway_handles_admin_global_models_locally_with_trusted_admin_principa
payload["models"].as_array().expect("models array")[0]["name"], payload["models"].as_array().expect("models array")[0]["name"],
"gpt-4.1" "gpt-4.1"
); );
assert_eq!(payload["models"][0]["provider_count"], 1); assert_eq!(payload["models"][0]["provider_count"], 3);
assert_eq!(payload["models"][0]["active_provider_count"], 1); assert_eq!(payload["models"][0]["active_provider_count"], 1);
assert_eq!(payload["models"][0]["usage_count"], 0); assert_eq!(payload["models"][0]["usage_count"], 7);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0); assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort(); gateway_handle.abort();
@@ -273,19 +291,35 @@ async fn gateway_handles_admin_global_model_detail_locally_with_trusted_admin_pr
}), }),
); );
let mut global_model = sample_admin_global_model("global-gpt-5", "gpt-5", "GPT 5");
global_model.usage_count = 7;
let mut anthropic_model = sample_admin_provider_model(
"model-anthropic-gpt5",
"provider-anthropic",
"global-gpt-5",
"gpt-5-anthropic",
);
anthropic_model.is_available = false;
let mut google_model = sample_admin_provider_model(
"model-google-gpt5",
"provider-google",
"global-gpt-5",
"gpt-5-google",
);
google_model.is_available = false;
let global_model_repository = Arc::new( let global_model_repository = Arc::new(
InMemoryGlobalModelReadRepository::seed(Vec::new()) InMemoryGlobalModelReadRepository::seed(Vec::new())
.with_admin_global_models(vec![sample_admin_global_model( .with_admin_global_models(vec![global_model])
"global-gpt-5", .with_admin_provider_models(vec![
"gpt-5", sample_admin_provider_model(
"GPT 5", "model-openai-gpt5",
)]) "provider-openai",
.with_admin_provider_models(vec![sample_admin_provider_model( "global-gpt-5",
"model-openai-gpt5", "gpt-5-upstream",
"provider-openai", ),
"global-gpt-5", anthropic_model,
"gpt-5-upstream", google_model,
)]), ]),
); );
let (upstream_url, upstream_handle) = start_server(upstream).await; let (upstream_url, upstream_handle) = start_server(upstream).await;
@@ -314,11 +348,11 @@ async fn gateway_handles_admin_global_model_detail_locally_with_trusted_admin_pr
assert_eq!(response.status(), StatusCode::OK); assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse"); let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["id"], "global-gpt-5"); assert_eq!(payload["id"], "global-gpt-5");
assert_eq!(payload["provider_count"], 1); assert_eq!(payload["provider_count"], 3);
assert_eq!(payload["active_provider_count"], 1); assert_eq!(payload["active_provider_count"], 1);
assert_eq!(payload["usage_count"], 0); assert_eq!(payload["usage_count"], 7);
assert_eq!(payload["total_models"], 1); assert_eq!(payload["total_models"], 3);
assert_eq!(payload["total_providers"], 1); assert_eq!(payload["total_providers"], 3);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0); assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort(); gateway_handle.abort();

View File

@@ -106,7 +106,9 @@ impl InMemoryGlobalModelReadRepository {
.len() as u64; .len() as u64;
let active_provider_count = items let active_provider_count = items
.iter() .iter()
.filter(|item| item.global_model_id == global_model_id && item.is_active) .filter(|item| {
item.global_model_id == global_model_id && item.is_active && item.is_available
})
.map(|item| item.provider_id.clone()) .map(|item| item.provider_id.clone())
.collect::<BTreeSet<_>>() .collect::<BTreeSet<_>>()
.len() as u64; .len() as u64;

View File

@@ -124,7 +124,7 @@ LEFT JOIN (
COUNT(DISTINCT m.provider_id)::bigint AS provider_count, COUNT(DISTINCT m.provider_id)::bigint AS provider_count,
COUNT( COUNT(
DISTINCT CASE DISTINCT CASE
WHEN m.is_active = TRUE AND p.is_active = TRUE THEN m.provider_id WHEN m.is_active = TRUE AND COALESCE(m.is_available, TRUE) = TRUE AND p.is_active = TRUE THEN m.provider_id
ELSE NULL ELSE NULL
END END
)::bigint AS active_provider_count )::bigint AS active_provider_count
@@ -411,7 +411,7 @@ LEFT JOIN (
COUNT(DISTINCT m.provider_id)::bigint AS provider_count, COUNT(DISTINCT m.provider_id)::bigint AS provider_count,
COUNT( COUNT(
DISTINCT CASE DISTINCT CASE
WHEN m.is_active = TRUE AND p.is_active = TRUE THEN m.provider_id WHEN m.is_active = TRUE AND COALESCE(m.is_available, TRUE) = TRUE AND p.is_active = TRUE THEN m.provider_id
ELSE NULL ELSE NULL
END END
)::bigint AS active_provider_count )::bigint AS active_provider_count
@@ -458,7 +458,7 @@ LEFT JOIN (
COUNT(DISTINCT m.provider_id)::bigint AS provider_count, COUNT(DISTINCT m.provider_id)::bigint AS provider_count,
COUNT( COUNT(
DISTINCT CASE DISTINCT CASE
WHEN m.is_active = TRUE AND p.is_active = TRUE THEN m.provider_id WHEN m.is_active = TRUE AND COALESCE(m.is_available, TRUE) = TRUE AND p.is_active = TRUE THEN m.provider_id
ELSE NULL ELSE NULL
END END
)::bigint AS active_provider_count )::bigint AS active_provider_count