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

feat(grok): add admin oauth and quota support

Browse Source
This commit is contained in:
mayrain
2026-05-16 21:15:38 +08:00
parent cbfe1d378f
commit 936e1ae37b
61 changed files with 4748 additions and 287 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/handlers/admin/provider/query/models/mod.rs
+341 -5
View File
@@ -17,6 +17,9 @@ use crate::ai_serving::{
};
use crate::clock::current_unix_ms;
use crate::execution_runtime;
use crate::handlers::admin::provider::shared::model_test_capabilities::{
admin_provider_model_supports_image_generation, admin_provider_model_test_capabilities_payload,
};
use crate::handlers::admin::request::{AdminAppState, AdminGatewayProviderTransportSnapshot};
use crate::handlers::shared::provider_pool::{
admin_provider_pool_config_from_config_value, read_admin_provider_pool_runtime_state,
@@ -100,6 +103,181 @@ struct ProviderQueryKeyFetchResult {
has_success: bool,
}
fn provider_query_model_id(model: &Value) -> Option<&str> {
model
.get("id")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
}
fn provider_query_grok_required_tier_rank(model_id: &str) -> Option<u8> {
match model_id.trim() {
"grok-4.20-0309-non-reasoning" | "grok-4.20-fast" | "grok-imagine-image-lite" => Some(0),
"grok-4.20-0309"
| "grok-4.20-0309-reasoning"
| "grok-4.20-0309-non-reasoning-super"
| "grok-4.20-0309-super"
| "grok-4.20-0309-reasoning-super"
| "grok-4.20-auto"
| "grok-4.20-expert"
| "grok-4.3-beta"
| "grok-imagine-image"
| "grok-imagine-image-pro"
| "grok-imagine-image-edit" => Some(1),
"grok-4.20-0309-non-reasoning-heavy"
| "grok-4.20-0309-heavy"
| "grok-4.20-0309-reasoning-heavy"
| "grok-4.20-multi-agent-0309"
| "grok-4.20-heavy" => Some(2),
_ => None,
}
}
fn provider_query_normalize_grok_pool_tier(value: Option<&str>) -> Option<&'static str> {
match value?.trim().to_ascii_lowercase().as_str() {
"basic" => Some("basic"),
"super" => Some("super"),
"heavy" => Some("heavy"),
_ => None,
}
}
fn provider_query_grok_pool_tier_rank(value: Option<&str>) -> u8 {
match provider_query_normalize_grok_pool_tier(value).unwrap_or("basic") {
"heavy" => 2,
"super" => 1,
_ => 0,
}
}
fn provider_query_grok_quota_string(quota: &Map<String, Value>, fields: &[&str]) -> Option<String> {
fields.iter().find_map(|field| {
quota
.get(*field)
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
})
}
fn provider_query_grok_window_limit(quota: &Map<String, Value>, model_name: &str) -> Option<f64> {
quota
.get("windows")
.and_then(Value::as_array)?
.iter()
.filter_map(Value::as_object)
.find(|window| {
window
.get("model")
.and_then(Value::as_str)
.is_some_and(|value| value.trim() == model_name)
})
.and_then(|window| window.get("limit_value"))
.and_then(Value::as_f64)
.filter(|value| value.is_finite() && *value > 0.0)
}
fn provider_query_grok_pool_tier_from_quota(quota: &Map<String, Value>) -> Option<&'static str> {
if let Some(tier) =
provider_query_grok_quota_string(quota, &["pool_tier", "tier", "plan_type", "plan"])
.and_then(|value| provider_query_normalize_grok_pool_tier(Some(&value)))
{
return Some(tier);
}
if let Some(auto_total) = provider_query_grok_window_limit(quota, "quota_auto") {
if (auto_total - 150.0).abs() < f64::EPSILON {
return Some("heavy");
}
if (auto_total - 50.0).abs() < f64::EPSILON {
return Some("super");
}
}
if let Some(fast_total) = provider_query_grok_window_limit(quota, "quota_fast") {
if (fast_total - 400.0).abs() < f64::EPSILON {
return Some("heavy");
}
if (fast_total - 140.0).abs() < f64::EPSILON {
return Some("super");
}
if (fast_total - 30.0).abs() < f64::EPSILON {
return Some("basic");
}
}
None
}
fn provider_query_grok_key_pool_tier(key: &StoredProviderCatalogKey) -> Option<&'static str> {
key.status_snapshot
.as_ref()
.and_then(Value::as_object)
.and_then(|snapshot| snapshot.get("quota"))
.and_then(Value::as_object)
.and_then(provider_query_grok_pool_tier_from_quota)
}
fn provider_query_filter_models_for_key(
provider: &StoredProviderCatalogProvider,
key: &StoredProviderCatalogKey,
models: Vec<Value>,
) -> Vec<Value> {
if !provider.provider_type.trim().eq_ignore_ascii_case("grok") {
return models;
}
let allowed_rank = provider_query_grok_pool_tier_rank(provider_query_grok_key_pool_tier(key));
models
.into_iter()
.filter(|model| {
provider_query_model_id(model)
.and_then(provider_query_grok_required_tier_rank)
.is_some_and(|required_rank| required_rank <= allowed_rank)
})
.collect()
}
fn provider_query_attach_model_test_capabilities(
provider: &StoredProviderCatalogProvider,
models: Vec<Value>,
) -> Vec<Value> {
models
.into_iter()
.map(|mut model| {
let Some(object) = model.as_object_mut() else {
return model;
};
let model_id = object
.get("id")
.and_then(Value::as_str)
.map(str::trim)
.unwrap_or_default()
.to_string();
let supports_image_generation = admin_provider_model_supports_image_generation(
&provider.provider_type,
&model_id,
object
.get("supports_image_generation")
.or_else(|| object.get("effective_supports_image_generation"))
.and_then(Value::as_bool)
.unwrap_or(false),
);
object.insert(
"model_test_capabilities".to_string(),
admin_provider_model_test_capabilities_payload(
&provider.provider_type,
&model_id,
supports_image_generation,
),
);
model
})
.collect()
}
fn provider_query_codex_preset_fallback(
provider: &StoredProviderCatalogProvider,
) -> Option<ProviderQueryKeyFetchResult> {
@@ -245,8 +423,9 @@ async fn provider_query_fetch_models_for_key(
if let Some(cached_models) =
provider_query_read_cached_models(state, &provider.id, &key.id).await
{
let models = provider_query_filter_models_for_key(provider, key, cached_models);
return Ok(ProviderQueryKeyFetchResult {
models: cached_models,
models,
error: None,
from_cache: true,
has_success: true,
@@ -257,8 +436,13 @@ async fn provider_query_fetch_models_for_key(
let selected_endpoints = selected_models_fetch_endpoints(endpoints, key);
if selected_endpoints.is_empty() {
if let Some(models) = preset_models_for_provider(&provider.provider_type) {
let models = provider_query_filter_models_for_key(
provider,
key,
aggregate_models_for_cache(&models),
);
return Ok(ProviderQueryKeyFetchResult {
models: aggregate_models_for_cache(&models),
models,
error: None,
from_cache: false,
has_success: true,
@@ -342,7 +526,7 @@ async fn provider_query_fetch_models_for_key(
}
Ok(ProviderQueryKeyFetchResult {
models: unique_models,
models: provider_query_filter_models_for_key(provider, key, unique_models),
error,
from_cache: false,
has_success: outcome.has_success,
@@ -397,11 +581,12 @@ pub(crate) async fn build_admin_provider_query_models_response(
force_refresh,
)
.await?;
let success = !result.models.is_empty();
let models = provider_query_attach_model_test_capabilities(&provider, result.models);
let success = !models.is_empty();
return Ok(Json(json!({
"success": success,
"data": {
"models": result.models,
"models": models,
"error": result.error,
"from_cache": result.from_cache,
},
@@ -429,6 +614,7 @@ pub(crate) async fn build_admin_provider_query_models_response(
{
if let Some(models) = provider_query_read_provider_cached_models(state, &provider.id).await
{
let models = provider_query_attach_model_test_capabilities(&provider, models);
return Ok(Json(json!({
"success": !models.is_empty(),
"data": {
@@ -504,6 +690,7 @@ pub(crate) async fn build_admin_provider_query_models_response(
if !success && error.is_none() {
error = Some(ADMIN_PROVIDER_QUERY_NO_MODELS_FROM_KEY_DETAIL.to_string());
}
let models = provider_query_attach_model_test_capabilities(&provider, models);
Ok(Json(json!({
"success": success,
@@ -519,3 +706,152 @@ pub(crate) async fn build_admin_provider_query_models_response(
}))
.into_response())
}
#[cfg(test)]
mod tests {
use super::*;
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
fn grok_provider() -> StoredProviderCatalogProvider {
let mut provider = StoredProviderCatalogProvider::new(
"provider-1".to_string(),
"Grok".to_string(),
None,
"grok".to_string(),
)
.expect("provider should build");
provider.provider_type = "grok".to_string();
provider
}
fn grok_key_with_quota(quota: Value) -> StoredProviderCatalogKey {
let mut key = StoredProviderCatalogKey::new(
"key-1".to_string(),
"provider-1".to_string(),
"key-1".to_string(),
"oauth".to_string(),
None,
true,
)
.expect("key should build");
key.status_snapshot = Some(json!({ "quota": quota }));
key
}
fn model(id: &str) -> Value {
json!({ "id": id })
}
fn filtered_ids(key: &StoredProviderCatalogKey) -> Vec<String> {
provider_query_filter_models_for_key(
&grok_provider(),
key,
vec![
model("grok-4.20-0309-non-reasoning"),
model("grok-4.20-auto"),
model("grok-4.20-heavy"),
model("grok-imagine-image-lite"),
model("grok-imagine-image"),
model("grok-imagine-image-edit"),
],
)
.into_iter()
.filter_map(|item| item.get("id").and_then(Value::as_str).map(str::to_string))
.collect()
}
#[test]
fn provider_query_grok_basic_tier_hides_super_and_heavy_models() {
let key = grok_key_with_quota(json!({ "pool_tier": "basic" }));
assert_eq!(
filtered_ids(&key),
["grok-4.20-0309-non-reasoning", "grok-imagine-image-lite"]
);
}
#[test]
fn provider_query_grok_super_tier_hides_heavy_models() {
let key = grok_key_with_quota(json!({ "plan_type": "super" }));
assert_eq!(
filtered_ids(&key),
[
"grok-4.20-0309-non-reasoning",
"grok-4.20-auto",
"grok-imagine-image-lite",
"grok-imagine-image",
"grok-imagine-image-edit"
]
);
}
#[test]
fn provider_query_grok_heavy_tier_keeps_full_non_video_catalog() {
let key = grok_key_with_quota(json!({ "pool_tier": "heavy" }));
assert_eq!(
filtered_ids(&key),
[
"grok-4.20-0309-non-reasoning",
"grok-4.20-auto",
"grok-4.20-heavy",
"grok-imagine-image-lite",
"grok-imagine-image",
"grok-imagine-image-edit"
]
);
}
#[test]
fn provider_query_grok_tier_falls_back_to_live_quota_windows() {
let key = grok_key_with_quota(json!({
"windows": [
{ "model": "quota_fast", "limit_value": 140.0 }
]
}));
assert_eq!(
filtered_ids(&key),
[
"grok-4.20-0309-non-reasoning",
"grok-4.20-auto",
"grok-imagine-image-lite",
"grok-imagine-image",
"grok-imagine-image-edit"
]
);
}
#[test]
fn provider_query_attaches_model_test_capabilities_to_models() {
let models = provider_query_attach_model_test_capabilities(
&grok_provider(),
vec![
model("grok-4.20-fast"),
model("grok-imagine-image"),
model("grok-imagine-image-edit"),
],
);
assert!(models[0]["model_test_capabilities"]["openai:image"].is_null());
assert_eq!(
models[1]["model_test_capabilities"]["openai:image"]["max_generation_count"],
json!(4)
);
assert_eq!(
models[1]["model_test_capabilities"]["openai:image"]["supports_generation"],
json!(true)
);
assert_eq!(
models[2]["model_test_capabilities"]["openai:image"]["supports_generation"],
json!(false)
);
assert_eq!(
models[2]["model_test_capabilities"]["openai:image"]["supports_edit"],
json!(true)
);
}
}