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

Merge remote-tracking branch 'origin/pr-493'

Browse Source
This commit is contained in:
fawney19
2026-05-18 18:01:43 +08:00
parent a8a87d2b41 89d08d9953
commit be939f7e63
4 changed files with 78 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
crates/aether-data/src/repository/provider_catalog/mysql.rs
View File

@@ -669,6 +669,11 @@ WHERE id = ?
"provider_api_keys.last_probe_increase_at",
)?)
.bind(optional_i64_from_u32(key.last_rpm_peak))
.bind(optional_i64_from_u64(
key.last_models_fetch_at_unix_secs,
"provider_api_keys.last_models_fetch_at",
)?)
.bind(&key.last_models_fetch_error)
.bind(updated_at)
.bind(&key.id)
.execute(&self.pool)
@@ -1208,6 +1213,8 @@ SET
utilization_samples = ?,
last_probe_increase_at = ?,
last_rpm_peak = ?,
last_models_fetch_at = ?,
last_models_fetch_error = ?,
updated_at = ?
WHERE id = ?
"#

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

@@ -1798,7 +1798,12 @@ SET
updated_at = CASE
WHEN $38::double precision IS NULL THEN NOW()
ELSE TO_TIMESTAMP($38::double precision)
END
END,
last_models_fetch_at = CASE
WHEN $42::double precision IS NULL THEN NULL
ELSE TO_TIMESTAMP($42::double precision)
END,
last_models_fetch_error = $43
WHERE id = $1
"#,
)
@@ -1846,6 +1851,8 @@ WHERE id = $1
.bind(key.expires_at_unix_secs.map(|value| value as f64))
.bind(&key.auth_type_by_format)
.bind(&key.allow_auth_channel_mismatch_formats)
.bind(key.last_models_fetch_at_unix_secs.map(|value| value as f64))
.bind(&key.last_models_fetch_error)
.execute(&self.pool)
.await
.map_postgres_err()?

22
crates/aether-data/src/repository/provider_catalog/sqlite.rs
View File

@@ -671,6 +671,11 @@ WHERE id = ?
"provider_api_keys.last_probe_increase_at",
)?)
.bind(optional_i64_from_u32(key.last_rpm_peak))
.bind(optional_i64_from_u64(
key.last_models_fetch_at_unix_secs,
"provider_api_keys.last_models_fetch_at",
)?)
.bind(&key.last_models_fetch_error)
.bind(updated_at)
.bind(&key.id)
.execute(&self.pool)
@@ -1210,6 +1215,8 @@ SET
utilization_samples = ?,
last_probe_increase_at = ?,
last_rpm_peak = ?,
last_models_fetch_at = ?,
last_models_fetch_error = ?,
updated_at = ?
WHERE id = ?
"#
@@ -1702,7 +1709,7 @@ mod tests {
assert_eq!(updated_endpoint.health_score, 0.5);
assert!(!updated_endpoint.is_active);
let key = StoredProviderCatalogKey::new(
let mut key = StoredProviderCatalogKey::new(
"key-write-1".to_string(),
"provider-write-1".to_string(),
"Default Key".to_string(),
@@ -1740,23 +1747,36 @@ mod tests {
Some(json!({"openai:chat":{"score":1}})),
Some(json!({"openai:chat":{"open":false}})),
);
key.last_models_fetch_at_unix_secs = Some(1_730_000_100);
key.last_models_fetch_error = Some("stale models fetch error".to_string());
let created_key = repository
.create_key(&key)
.await
.expect("key should create");
assert_eq!(created_key.concurrent_limit, Some(3));
assert_eq!(created_key.total_tokens, 1234);
assert_eq!(
created_key.last_models_fetch_error.as_deref(),
Some("stale models fetch error")
);
let mut updated_key = created_key.clone();
updated_key.name = "Updated Key".to_string();
updated_key.is_active = false;
updated_key.upstream_metadata = Some(json!({"models":["gpt-4.1"]}));
updated_key.last_models_fetch_at_unix_secs = Some(1_730_000_200);
updated_key.last_models_fetch_error = None;
let updated_key = repository
.update_key(&updated_key)
.await
.expect("key should update");
assert_eq!(updated_key.name, "Updated Key");
assert!(!updated_key.is_active);
assert_eq!(
updated_key.last_models_fetch_at_unix_secs,
Some(1_730_000_200)
);
assert_eq!(updated_key.last_models_fetch_error, None);
assert!(repository
.update_key_upstream_metadata(

57
crates/aether-model-fetch/src/strategy.rs
View File

@@ -28,6 +28,7 @@ const ANTIGRAVITY_DAILY_BASE_URL: &str = "https://daily-cloudcode-pa.googleapis.
const ANTIGRAVITY_PROD_BASE_URL: &str = "https://cloudcode-pa.googleapis.com";
const ANTIGRAVITY_BLOCKED_MODELS: &[&str] = &["chat_23310", "chat_20706"];
const VERTEX_API_BASE_URL: &str = "https://aiplatform.googleapis.com";
const VERTEX_MODEL_GARDEN_LIST_API_VERSION: &str = "v1beta1";
const VERTEX_PAGE_SIZE: &str = "100";
const VERTEX_MAX_PAGES: usize = 20;
const GOOGLE_OAUTH_TOKEN_URL: &str = "https://oauth2.googleapis.com/token";
@@ -929,14 +930,7 @@ fn iter_vertex_base_urls(transports: &[GatewayProviderTransportSnapshot]) -> Vec
}
fn build_vertex_google_list_url(base_url: &str, api_key: &str, page_token: Option<&str>) -> String {
let path = if base_url.trim_end_matches('/').ends_with("/v1")
|| base_url.trim_end_matches('/').ends_with("/v1beta")
{
"/publishers/google/models"
} else {
"/v1/publishers/google/models"
};
let url = build_simple_path_url(base_url, path);
let url = build_vertex_model_garden_list_url(base_url, "google");
let mut url = append_query_param(url, "key", api_key);
url = append_query_param(url, "pageSize", VERTEX_PAGE_SIZE);
if let Some(page_token) = page_token {
@@ -947,14 +941,12 @@ fn build_vertex_google_list_url(base_url: &str, api_key: &str, page_token: Optio
fn build_vertex_service_account_list_url(
base_url: &str,
project_id: &str,
region: &str,
_project_id: &str,
_region: &str,
publisher: &str,
page_token: Option<&str>,
) -> String {
let path =
format!("/v1/projects/{project_id}/locations/{region}/publishers/{publisher}/models");
let mut url = build_simple_path_url(base_url, &path);
let mut url = build_vertex_model_garden_list_url(base_url, publisher);
url = append_query_param(url, "pageSize", VERTEX_PAGE_SIZE);
if let Some(page_token) = page_token {
url = append_query_param(url, "pageToken", page_token);
@@ -962,6 +954,16 @@ fn build_vertex_service_account_list_url(
url
}
fn build_vertex_model_garden_list_url(base_url: &str, publisher: &str) -> String {
let trimmed_base = base_url.trim().trim_end_matches('/');
let unversioned_base = trimmed_base
.strip_suffix("/v1beta1")
.or_else(|| trimmed_base.strip_suffix("/v1"))
.unwrap_or(trimmed_base);
let path = format!("/{VERTEX_MODEL_GARDEN_LIST_API_VERSION}/publishers/{publisher}/models");
build_simple_path_url(unversioned_base, &path)
}
fn build_simple_path_url(base_url: &str, path: &str) -> String {
format!("{}{}", base_url.trim().trim_end_matches('/'), path.trim())
}
@@ -1304,7 +1306,10 @@ mod tests {
use async_trait::async_trait;
use serde_json::{json, Value};
use super::{select_model_fetch_strategy, ModelFetchStrategy, ModelFetchStrategyKind};
use super::{
build_vertex_google_list_url, build_vertex_service_account_list_url,
select_model_fetch_strategy, ModelFetchStrategy, ModelFetchStrategyKind,
};
use crate::fetch_models_from_transports;
use crate::transport::ModelFetchTransportRuntime;
@@ -1498,7 +1503,7 @@ mod tests {
let urls = executed_urls.lock().expect("executed_urls lock");
assert_eq!(
urls.as_slice(),
&["https://aiplatform.googleapis.com/v1/publishers/google/models?key=vertex-secret&pageSize=100"]
&["https://aiplatform.googleapis.com/v1beta1/publishers/google/models?key=vertex-secret&pageSize=100"]
);
assert_eq!(outcome.fetched_model_ids, vec!["gemini-3.1-pro-preview"]);
assert_eq!(outcome.cached_models.len(), 1);
@@ -1508,6 +1513,28 @@ mod tests {
);
}
#[test]
fn vertex_model_fetch_uses_model_garden_list_endpoint() {
assert_eq!(
build_vertex_google_list_url(
"https://aiplatform.googleapis.com/v1",
"vertex-secret",
None,
),
"https://aiplatform.googleapis.com/v1beta1/publishers/google/models?key=vertex-secret&pageSize=100"
);
assert_eq!(
build_vertex_service_account_list_url(
"https://aiplatform.googleapis.com",
"project-1",
"global",
"google",
Some("page-2"),
),
"https://aiplatform.googleapis.com/v1beta1/publishers/google/models?pageSize=100&pageToken=page-2"
);
}
#[tokio::test]
async fn codex_transport_fetches_upstream_models_instead_of_preset_catalog() {
let executed_urls = Arc::new(Mutex::new(Vec::new()));