mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-13 04:29:50 +08:00
Merge branch 'fawney19:main' into main
This commit is contained in:
@@ -334,7 +334,7 @@ impl GatewayDataState {
|
|||||||
encrypted_auth_config: Option<&str>,
|
encrypted_auth_config: Option<&str>,
|
||||||
expires_at_unix_secs: Option<u64>,
|
expires_at_unix_secs: Option<u64>,
|
||||||
) -> Result<bool, DataLayerError> {
|
) -> Result<bool, DataLayerError> {
|
||||||
match &self.provider_catalog_writer {
|
let updated = match &self.provider_catalog_writer {
|
||||||
Some(repository) => {
|
Some(repository) => {
|
||||||
repository
|
repository
|
||||||
.update_key_oauth_credentials(
|
.update_key_oauth_credentials(
|
||||||
@@ -346,17 +346,25 @@ impl GatewayDataState {
|
|||||||
.await
|
.await
|
||||||
}
|
}
|
||||||
None => Ok(false),
|
None => Ok(false),
|
||||||
|
}?;
|
||||||
|
if updated {
|
||||||
|
self.clear_provider_catalog_cache();
|
||||||
}
|
}
|
||||||
|
Ok(updated)
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(crate) async fn create_provider_catalog_key(
|
pub(crate) async fn create_provider_catalog_key(
|
||||||
&self,
|
&self,
|
||||||
key: &StoredProviderCatalogKey,
|
key: &StoredProviderCatalogKey,
|
||||||
) -> Result<Option<StoredProviderCatalogKey>, DataLayerError> {
|
) -> Result<Option<StoredProviderCatalogKey>, DataLayerError> {
|
||||||
match &self.provider_catalog_writer {
|
let created = match &self.provider_catalog_writer {
|
||||||
Some(repository) => repository.create_key(key).await.map(Some),
|
Some(repository) => repository.create_key(key).await.map(Some),
|
||||||
None => Ok(None),
|
None => Ok(None),
|
||||||
|
}?;
|
||||||
|
if created.is_some() {
|
||||||
|
self.clear_provider_catalog_cache();
|
||||||
}
|
}
|
||||||
|
Ok(created)
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(crate) async fn create_provider_catalog_provider(
|
pub(crate) async fn create_provider_catalog_provider(
|
||||||
@@ -364,33 +372,45 @@ impl GatewayDataState {
|
|||||||
provider: &StoredProviderCatalogProvider,
|
provider: &StoredProviderCatalogProvider,
|
||||||
shift_existing_priorities_from: Option<i32>,
|
shift_existing_priorities_from: Option<i32>,
|
||||||
) -> Result<Option<StoredProviderCatalogProvider>, DataLayerError> {
|
) -> Result<Option<StoredProviderCatalogProvider>, DataLayerError> {
|
||||||
match &self.provider_catalog_writer {
|
let created = match &self.provider_catalog_writer {
|
||||||
Some(repository) => repository
|
Some(repository) => repository
|
||||||
.create_provider(provider, shift_existing_priorities_from)
|
.create_provider(provider, shift_existing_priorities_from)
|
||||||
.await
|
.await
|
||||||
.map(Some),
|
.map(Some),
|
||||||
None => Ok(None),
|
None => Ok(None),
|
||||||
|
}?;
|
||||||
|
if created.is_some() {
|
||||||
|
self.clear_provider_catalog_cache();
|
||||||
}
|
}
|
||||||
|
Ok(created)
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(crate) async fn update_provider_catalog_provider(
|
pub(crate) async fn update_provider_catalog_provider(
|
||||||
&self,
|
&self,
|
||||||
provider: &StoredProviderCatalogProvider,
|
provider: &StoredProviderCatalogProvider,
|
||||||
) -> Result<Option<StoredProviderCatalogProvider>, DataLayerError> {
|
) -> Result<Option<StoredProviderCatalogProvider>, DataLayerError> {
|
||||||
match &self.provider_catalog_writer {
|
let updated = match &self.provider_catalog_writer {
|
||||||
Some(repository) => repository.update_provider(provider).await.map(Some),
|
Some(repository) => repository.update_provider(provider).await.map(Some),
|
||||||
None => Ok(None),
|
None => Ok(None),
|
||||||
|
}?;
|
||||||
|
if updated.is_some() {
|
||||||
|
self.clear_provider_catalog_cache();
|
||||||
}
|
}
|
||||||
|
Ok(updated)
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(crate) async fn delete_provider_catalog_provider(
|
pub(crate) async fn delete_provider_catalog_provider(
|
||||||
&self,
|
&self,
|
||||||
provider_id: &str,
|
provider_id: &str,
|
||||||
) -> Result<bool, DataLayerError> {
|
) -> Result<bool, DataLayerError> {
|
||||||
match &self.provider_catalog_writer {
|
let deleted = match &self.provider_catalog_writer {
|
||||||
Some(repository) => repository.delete_provider(provider_id).await,
|
Some(repository) => repository.delete_provider(provider_id).await,
|
||||||
None => Ok(false),
|
None => Ok(false),
|
||||||
|
}?;
|
||||||
|
if deleted {
|
||||||
|
self.clear_provider_catalog_cache();
|
||||||
}
|
}
|
||||||
|
Ok(deleted)
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(crate) async fn cleanup_deleted_provider_catalog_refs(
|
pub(crate) async fn cleanup_deleted_provider_catalog_refs(
|
||||||
@@ -399,54 +419,74 @@ impl GatewayDataState {
|
|||||||
endpoint_ids: &[String],
|
endpoint_ids: &[String],
|
||||||
key_ids: &[String],
|
key_ids: &[String],
|
||||||
) -> Result<(), DataLayerError> {
|
) -> Result<(), DataLayerError> {
|
||||||
match &self.provider_catalog_writer {
|
let cleaned = match &self.provider_catalog_writer {
|
||||||
Some(repository) => {
|
Some(repository) => {
|
||||||
repository
|
repository
|
||||||
.cleanup_deleted_provider_refs(provider_id, endpoint_ids, key_ids)
|
.cleanup_deleted_provider_refs(provider_id, endpoint_ids, key_ids)
|
||||||
.await
|
.await
|
||||||
}
|
}
|
||||||
None => Ok(()),
|
None => Ok(()),
|
||||||
|
};
|
||||||
|
if !endpoint_ids.is_empty() || !key_ids.is_empty() {
|
||||||
|
self.clear_provider_catalog_cache();
|
||||||
}
|
}
|
||||||
|
cleaned
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(crate) async fn create_provider_catalog_endpoint(
|
pub(crate) async fn create_provider_catalog_endpoint(
|
||||||
&self,
|
&self,
|
||||||
endpoint: &StoredProviderCatalogEndpoint,
|
endpoint: &StoredProviderCatalogEndpoint,
|
||||||
) -> Result<Option<StoredProviderCatalogEndpoint>, DataLayerError> {
|
) -> Result<Option<StoredProviderCatalogEndpoint>, DataLayerError> {
|
||||||
match &self.provider_catalog_writer {
|
let created = match &self.provider_catalog_writer {
|
||||||
Some(repository) => repository.create_endpoint(endpoint).await.map(Some),
|
Some(repository) => repository.create_endpoint(endpoint).await.map(Some),
|
||||||
None => Ok(None),
|
None => Ok(None),
|
||||||
|
}?;
|
||||||
|
if created.is_some() {
|
||||||
|
self.clear_provider_catalog_cache();
|
||||||
}
|
}
|
||||||
|
Ok(created)
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(crate) async fn update_provider_catalog_endpoint(
|
pub(crate) async fn update_provider_catalog_endpoint(
|
||||||
&self,
|
&self,
|
||||||
endpoint: &StoredProviderCatalogEndpoint,
|
endpoint: &StoredProviderCatalogEndpoint,
|
||||||
) -> Result<Option<StoredProviderCatalogEndpoint>, DataLayerError> {
|
) -> Result<Option<StoredProviderCatalogEndpoint>, DataLayerError> {
|
||||||
match &self.provider_catalog_writer {
|
let updated = match &self.provider_catalog_writer {
|
||||||
Some(repository) => repository.update_endpoint(endpoint).await.map(Some),
|
Some(repository) => repository.update_endpoint(endpoint).await.map(Some),
|
||||||
None => Ok(None),
|
None => Ok(None),
|
||||||
|
}?;
|
||||||
|
if updated.is_some() {
|
||||||
|
self.clear_provider_catalog_cache();
|
||||||
}
|
}
|
||||||
|
Ok(updated)
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(crate) async fn delete_provider_catalog_endpoint(
|
pub(crate) async fn delete_provider_catalog_endpoint(
|
||||||
&self,
|
&self,
|
||||||
endpoint_id: &str,
|
endpoint_id: &str,
|
||||||
) -> Result<bool, DataLayerError> {
|
) -> Result<bool, DataLayerError> {
|
||||||
match &self.provider_catalog_writer {
|
let deleted = match &self.provider_catalog_writer {
|
||||||
Some(repository) => repository.delete_endpoint(endpoint_id).await,
|
Some(repository) => repository.delete_endpoint(endpoint_id).await,
|
||||||
None => Ok(false),
|
None => Ok(false),
|
||||||
|
}?;
|
||||||
|
if deleted {
|
||||||
|
self.clear_provider_catalog_cache();
|
||||||
}
|
}
|
||||||
|
Ok(deleted)
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(crate) async fn update_provider_catalog_key(
|
pub(crate) async fn update_provider_catalog_key(
|
||||||
&self,
|
&self,
|
||||||
key: &StoredProviderCatalogKey,
|
key: &StoredProviderCatalogKey,
|
||||||
) -> Result<Option<StoredProviderCatalogKey>, DataLayerError> {
|
) -> Result<Option<StoredProviderCatalogKey>, DataLayerError> {
|
||||||
match &self.provider_catalog_writer {
|
let updated = match &self.provider_catalog_writer {
|
||||||
Some(repository) => repository.update_key(key).await.map(Some),
|
Some(repository) => repository.update_key(key).await.map(Some),
|
||||||
None => Ok(None),
|
None => Ok(None),
|
||||||
|
}?;
|
||||||
|
if updated.is_some() {
|
||||||
|
self.clear_provider_catalog_cache();
|
||||||
}
|
}
|
||||||
|
Ok(updated)
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(crate) async fn update_provider_catalog_key_upstream_metadata(
|
pub(crate) async fn update_provider_catalog_key_upstream_metadata(
|
||||||
@@ -455,34 +495,46 @@ impl GatewayDataState {
|
|||||||
upstream_metadata: Option<&serde_json::Value>,
|
upstream_metadata: Option<&serde_json::Value>,
|
||||||
updated_at_unix_secs: Option<u64>,
|
updated_at_unix_secs: Option<u64>,
|
||||||
) -> Result<bool, DataLayerError> {
|
) -> Result<bool, DataLayerError> {
|
||||||
match &self.provider_catalog_writer {
|
let updated = match &self.provider_catalog_writer {
|
||||||
Some(repository) => {
|
Some(repository) => {
|
||||||
repository
|
repository
|
||||||
.update_key_upstream_metadata(key_id, upstream_metadata, updated_at_unix_secs)
|
.update_key_upstream_metadata(key_id, upstream_metadata, updated_at_unix_secs)
|
||||||
.await
|
.await
|
||||||
}
|
}
|
||||||
None => Ok(false),
|
None => Ok(false),
|
||||||
|
}?;
|
||||||
|
if updated {
|
||||||
|
self.clear_provider_catalog_cache();
|
||||||
}
|
}
|
||||||
|
Ok(updated)
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(crate) async fn delete_provider_catalog_key(
|
pub(crate) async fn delete_provider_catalog_key(
|
||||||
&self,
|
&self,
|
||||||
key_id: &str,
|
key_id: &str,
|
||||||
) -> Result<bool, DataLayerError> {
|
) -> Result<bool, DataLayerError> {
|
||||||
match &self.provider_catalog_writer {
|
let deleted = match &self.provider_catalog_writer {
|
||||||
Some(repository) => repository.delete_key(key_id).await,
|
Some(repository) => repository.delete_key(key_id).await,
|
||||||
None => Ok(false),
|
None => Ok(false),
|
||||||
|
}?;
|
||||||
|
if deleted {
|
||||||
|
self.clear_provider_catalog_cache();
|
||||||
}
|
}
|
||||||
|
Ok(deleted)
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(crate) async fn clear_provider_catalog_key_oauth_invalid_marker(
|
pub(crate) async fn clear_provider_catalog_key_oauth_invalid_marker(
|
||||||
&self,
|
&self,
|
||||||
key_id: &str,
|
key_id: &str,
|
||||||
) -> Result<bool, DataLayerError> {
|
) -> Result<bool, DataLayerError> {
|
||||||
match &self.provider_catalog_writer {
|
let updated = match &self.provider_catalog_writer {
|
||||||
Some(repository) => repository.clear_key_oauth_invalid_marker(key_id).await,
|
Some(repository) => repository.clear_key_oauth_invalid_marker(key_id).await,
|
||||||
None => Ok(false),
|
None => Ok(false),
|
||||||
|
}?;
|
||||||
|
if updated {
|
||||||
|
self.clear_provider_catalog_cache();
|
||||||
}
|
}
|
||||||
|
Ok(updated)
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(crate) async fn update_provider_catalog_key_health_state(
|
pub(crate) async fn update_provider_catalog_key_health_state(
|
||||||
@@ -492,7 +544,7 @@ impl GatewayDataState {
|
|||||||
health_by_format: Option<&serde_json::Value>,
|
health_by_format: Option<&serde_json::Value>,
|
||||||
circuit_breaker_by_format: Option<&serde_json::Value>,
|
circuit_breaker_by_format: Option<&serde_json::Value>,
|
||||||
) -> Result<bool, DataLayerError> {
|
) -> Result<bool, DataLayerError> {
|
||||||
match &self.provider_catalog_writer {
|
let updated = match &self.provider_catalog_writer {
|
||||||
Some(repository) => {
|
Some(repository) => {
|
||||||
repository
|
repository
|
||||||
.update_key_health_state(
|
.update_key_health_state(
|
||||||
@@ -504,6 +556,10 @@ impl GatewayDataState {
|
|||||||
.await
|
.await
|
||||||
}
|
}
|
||||||
None => Ok(false),
|
None => Ok(false),
|
||||||
|
}?;
|
||||||
|
if updated {
|
||||||
|
self.clear_provider_catalog_cache();
|
||||||
}
|
}
|
||||||
|
Ok(updated)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1,5 +1,6 @@
|
|||||||
use aether_data::{DataBackends, DataLayerError, DatabaseDriver};
|
use aether_data::{DataBackends, DataLayerError, DatabaseDriver};
|
||||||
use aether_data_contracts::repository::candidate_selection::MinimalCandidateSelectionReadRepository;
|
use aether_data_contracts::repository::candidate_selection::MinimalCandidateSelectionReadRepository;
|
||||||
|
use aether_data_contracts::repository::provider_catalog::ProviderCatalogReadRepository;
|
||||||
use aether_runtime_state::RuntimeQueueStore;
|
use aether_runtime_state::RuntimeQueueStore;
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
|
|
||||||
@@ -99,7 +100,11 @@ impl GatewayDataState {
|
|||||||
let request_candidate_reader = backends.read().request_candidates();
|
let request_candidate_reader = backends.read().request_candidates();
|
||||||
let request_candidate_writer = backends.write().request_candidates();
|
let request_candidate_writer = backends.write().request_candidates();
|
||||||
let gemini_file_mapping_writer = backends.write().gemini_file_mappings();
|
let gemini_file_mapping_writer = backends.write().gemini_file_mappings();
|
||||||
let provider_catalog_reader = backends.read().provider_catalog();
|
let provider_catalog_reader = backends.read().provider_catalog().map(|repository| {
|
||||||
|
Arc::new(
|
||||||
|
super::provider_catalog_cache::CachedProviderCatalogReadRepository::new(repository),
|
||||||
|
) as Arc<dyn ProviderCatalogReadRepository>
|
||||||
|
});
|
||||||
let provider_catalog_writer = backends.write().provider_catalog();
|
let provider_catalog_writer = backends.write().provider_catalog();
|
||||||
let pool_score_reader = backends.read().pool_scores();
|
let pool_score_reader = backends.read().pool_scores();
|
||||||
let pool_score_writer = backends.write().pool_scores();
|
let pool_score_writer = backends.write().pool_scores();
|
||||||
@@ -276,6 +281,12 @@ impl GatewayDataState {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub(crate) fn clear_provider_catalog_cache(&self) {
|
||||||
|
if let Some(repository) = &self.provider_catalog_reader {
|
||||||
|
repository.clear_local_cache();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
pub(crate) fn has_request_candidate_reader(&self) -> bool {
|
pub(crate) fn has_request_candidate_reader(&self) -> bool {
|
||||||
self.request_candidate_reader.is_some()
|
self.request_candidate_reader.is_some()
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -324,6 +324,7 @@ mod core;
|
|||||||
mod integrations;
|
mod integrations;
|
||||||
mod models;
|
mod models;
|
||||||
mod pool_scores;
|
mod pool_scores;
|
||||||
|
mod provider_catalog_cache;
|
||||||
mod referrals;
|
mod referrals;
|
||||||
mod routing_profiles;
|
mod routing_profiles;
|
||||||
mod runtime;
|
mod runtime;
|
||||||
|
|||||||
@@ -0,0 +1,360 @@
|
|||||||
|
use std::collections::HashMap;
|
||||||
|
use std::future::Future;
|
||||||
|
use std::sync::atomic::{AtomicU64, Ordering};
|
||||||
|
use std::sync::{Arc, Mutex};
|
||||||
|
use std::time::Duration;
|
||||||
|
|
||||||
|
use aether_cache::ExpiringMap;
|
||||||
|
use aether_data::DataLayerError;
|
||||||
|
use aether_data_contracts::repository::provider_catalog::{
|
||||||
|
ProviderCatalogKeyListQuery, ProviderCatalogReadRepository, StoredProviderCatalogEndpoint,
|
||||||
|
StoredProviderCatalogKey, StoredProviderCatalogKeyMaintenanceSummary,
|
||||||
|
StoredProviderCatalogKeyPage, StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
|
||||||
|
};
|
||||||
|
use async_trait::async_trait;
|
||||||
|
use tokio::sync::Notify;
|
||||||
|
|
||||||
|
const PROVIDER_CATALOG_CACHE_TTL: Duration = Duration::from_secs(5);
|
||||||
|
const PROVIDER_CATALOG_CACHE_MAX_ENTRIES: usize = 1024;
|
||||||
|
|
||||||
|
pub(super) struct CachedProviderCatalogReadRepository {
|
||||||
|
inner: Arc<dyn ProviderCatalogReadRepository>,
|
||||||
|
entries: ExpiringMap<ProviderCatalogCacheKey, ProviderCatalogCacheValue>,
|
||||||
|
inflight: Mutex<HashMap<ProviderCatalogCacheKey, u64>>,
|
||||||
|
inflight_notify: Notify,
|
||||||
|
next_inflight_token: AtomicU64,
|
||||||
|
epoch: AtomicU64,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl CachedProviderCatalogReadRepository {
|
||||||
|
pub(super) fn new(inner: Arc<dyn ProviderCatalogReadRepository>) -> Self {
|
||||||
|
Self {
|
||||||
|
inner,
|
||||||
|
entries: ExpiringMap::new(),
|
||||||
|
inflight: Mutex::new(HashMap::new()),
|
||||||
|
inflight_notify: Notify::new(),
|
||||||
|
next_inflight_token: AtomicU64::new(1),
|
||||||
|
epoch: AtomicU64::new(0),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn get_or_load<F, Fut>(
|
||||||
|
&self,
|
||||||
|
key: ProviderCatalogCacheKey,
|
||||||
|
load: F,
|
||||||
|
) -> Result<ProviderCatalogCacheValue, DataLayerError>
|
||||||
|
where
|
||||||
|
F: Fn() -> Fut,
|
||||||
|
Fut: Future<Output = Result<ProviderCatalogCacheValue, DataLayerError>>,
|
||||||
|
{
|
||||||
|
if let Some(value) = self.entries.get_fresh(&key, PROVIDER_CATALOG_CACHE_TTL) {
|
||||||
|
return Ok(value);
|
||||||
|
}
|
||||||
|
|
||||||
|
loop {
|
||||||
|
let notified = self.inflight_notify.notified();
|
||||||
|
match self.register_inflight(&key) {
|
||||||
|
InflightRegistration::Bypass => return load().await,
|
||||||
|
InflightRegistration::Follower => {
|
||||||
|
notified.await;
|
||||||
|
if let Some(value) = self.entries.get_fresh(&key, PROVIDER_CATALOG_CACHE_TTL) {
|
||||||
|
return Ok(value);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
InflightRegistration::Leader(token) => {
|
||||||
|
let mut guard = InflightGuard::new(self, key.clone(), token);
|
||||||
|
let load_epoch = self.epoch.load(Ordering::Acquire);
|
||||||
|
let result = load().await;
|
||||||
|
if let Ok(value) = &result {
|
||||||
|
if load_epoch == self.epoch.load(Ordering::Acquire) {
|
||||||
|
self.entries.insert(
|
||||||
|
key.clone(),
|
||||||
|
value.clone(),
|
||||||
|
PROVIDER_CATALOG_CACHE_TTL,
|
||||||
|
PROVIDER_CATALOG_CACHE_MAX_ENTRIES,
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
guard.finish();
|
||||||
|
return result;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn register_inflight(&self, key: &ProviderCatalogCacheKey) -> InflightRegistration {
|
||||||
|
match self.inflight.lock() {
|
||||||
|
Ok(mut inflight) => {
|
||||||
|
if inflight.contains_key(key) {
|
||||||
|
return InflightRegistration::Follower;
|
||||||
|
}
|
||||||
|
let token = self.next_inflight_token.fetch_add(1, Ordering::AcqRel);
|
||||||
|
inflight.insert(key.clone(), token);
|
||||||
|
InflightRegistration::Leader(token)
|
||||||
|
}
|
||||||
|
Err(_) => InflightRegistration::Bypass,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn finish_inflight(&self, key: &ProviderCatalogCacheKey, token: u64) {
|
||||||
|
let mut removed = false;
|
||||||
|
if let Ok(mut inflight) = self.inflight.lock() {
|
||||||
|
if inflight.get(key).copied() == Some(token) {
|
||||||
|
inflight.remove(key);
|
||||||
|
removed = true;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if removed {
|
||||||
|
self.inflight_notify.notify_waiters();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn clear(&self) {
|
||||||
|
self.epoch.fetch_add(1, Ordering::AcqRel);
|
||||||
|
self.entries.clear();
|
||||||
|
let mut cleared_inflight = false;
|
||||||
|
if let Ok(mut inflight) = self.inflight.lock() {
|
||||||
|
cleared_inflight = !inflight.is_empty();
|
||||||
|
inflight.clear();
|
||||||
|
}
|
||||||
|
if cleared_inflight {
|
||||||
|
self.inflight_notify.notify_waiters();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl ProviderCatalogReadRepository for CachedProviderCatalogReadRepository {
|
||||||
|
fn clear_local_cache(&self) {
|
||||||
|
self.clear();
|
||||||
|
self.inner.clear_local_cache();
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn list_providers(
|
||||||
|
&self,
|
||||||
|
active_only: bool,
|
||||||
|
) -> Result<Vec<StoredProviderCatalogProvider>, DataLayerError> {
|
||||||
|
match self
|
||||||
|
.get_or_load(
|
||||||
|
ProviderCatalogCacheKey::Providers { active_only },
|
||||||
|
|| async move {
|
||||||
|
self.inner
|
||||||
|
.list_providers(active_only)
|
||||||
|
.await
|
||||||
|
.map(ProviderCatalogCacheValue::Providers)
|
||||||
|
},
|
||||||
|
)
|
||||||
|
.await?
|
||||||
|
{
|
||||||
|
ProviderCatalogCacheValue::Providers(items) => Ok(items),
|
||||||
|
_ => Ok(Vec::new()),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn list_providers_by_ids(
|
||||||
|
&self,
|
||||||
|
provider_ids: &[String],
|
||||||
|
) -> Result<Vec<StoredProviderCatalogProvider>, DataLayerError> {
|
||||||
|
let key = ProviderCatalogCacheKey::ProvidersByIds(normalize_ids(provider_ids));
|
||||||
|
match self
|
||||||
|
.get_or_load(key, || async move {
|
||||||
|
self.inner
|
||||||
|
.list_providers_by_ids(provider_ids)
|
||||||
|
.await
|
||||||
|
.map(ProviderCatalogCacheValue::Providers)
|
||||||
|
})
|
||||||
|
.await?
|
||||||
|
{
|
||||||
|
ProviderCatalogCacheValue::Providers(items) => Ok(items),
|
||||||
|
_ => Ok(Vec::new()),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn list_endpoints_by_ids(
|
||||||
|
&self,
|
||||||
|
endpoint_ids: &[String],
|
||||||
|
) -> Result<Vec<StoredProviderCatalogEndpoint>, DataLayerError> {
|
||||||
|
self.inner.list_endpoints_by_ids(endpoint_ids).await
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn list_endpoints_by_provider_ids(
|
||||||
|
&self,
|
||||||
|
provider_ids: &[String],
|
||||||
|
) -> Result<Vec<StoredProviderCatalogEndpoint>, DataLayerError> {
|
||||||
|
let key = ProviderCatalogCacheKey::EndpointsByProviderIds(normalize_ids(provider_ids));
|
||||||
|
match self
|
||||||
|
.get_or_load(key, || async move {
|
||||||
|
self.inner
|
||||||
|
.list_endpoints_by_provider_ids(provider_ids)
|
||||||
|
.await
|
||||||
|
.map(ProviderCatalogCacheValue::Endpoints)
|
||||||
|
})
|
||||||
|
.await?
|
||||||
|
{
|
||||||
|
ProviderCatalogCacheValue::Endpoints(items) => Ok(items),
|
||||||
|
_ => Ok(Vec::new()),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn list_keys_by_ids(
|
||||||
|
&self,
|
||||||
|
key_ids: &[String],
|
||||||
|
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
|
||||||
|
self.inner.list_keys_by_ids(key_ids).await
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn list_keys_by_provider_ids(
|
||||||
|
&self,
|
||||||
|
provider_ids: &[String],
|
||||||
|
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
|
||||||
|
let key = ProviderCatalogCacheKey::KeysByProviderIds(normalize_ids(provider_ids));
|
||||||
|
match self
|
||||||
|
.get_or_load(key, || async move {
|
||||||
|
self.inner
|
||||||
|
.list_keys_by_provider_ids(provider_ids)
|
||||||
|
.await
|
||||||
|
.map(ProviderCatalogCacheValue::Keys)
|
||||||
|
})
|
||||||
|
.await?
|
||||||
|
{
|
||||||
|
ProviderCatalogCacheValue::Keys(items) => Ok(items),
|
||||||
|
_ => Ok(Vec::new()),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn list_key_summaries_by_provider_ids(
|
||||||
|
&self,
|
||||||
|
provider_ids: &[String],
|
||||||
|
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
|
||||||
|
let key = ProviderCatalogCacheKey::KeySummariesByProviderIds(normalize_ids(provider_ids));
|
||||||
|
match self
|
||||||
|
.get_or_load(key, || async move {
|
||||||
|
self.inner
|
||||||
|
.list_key_summaries_by_provider_ids(provider_ids)
|
||||||
|
.await
|
||||||
|
.map(ProviderCatalogCacheValue::Keys)
|
||||||
|
})
|
||||||
|
.await?
|
||||||
|
{
|
||||||
|
ProviderCatalogCacheValue::Keys(items) => Ok(items),
|
||||||
|
_ => Ok(Vec::new()),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn list_key_maintenance_summaries_by_provider_ids(
|
||||||
|
&self,
|
||||||
|
provider_ids: &[String],
|
||||||
|
) -> Result<Vec<StoredProviderCatalogKeyMaintenanceSummary>, DataLayerError> {
|
||||||
|
let key = ProviderCatalogCacheKey::KeyMaintenanceSummariesByProviderIds(normalize_ids(
|
||||||
|
provider_ids,
|
||||||
|
));
|
||||||
|
match self
|
||||||
|
.get_or_load(key, || async move {
|
||||||
|
self.inner
|
||||||
|
.list_key_maintenance_summaries_by_provider_ids(provider_ids)
|
||||||
|
.await
|
||||||
|
.map(ProviderCatalogCacheValue::KeyMaintenanceSummaries)
|
||||||
|
})
|
||||||
|
.await?
|
||||||
|
{
|
||||||
|
ProviderCatalogCacheValue::KeyMaintenanceSummaries(items) => Ok(items),
|
||||||
|
_ => Ok(Vec::new()),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn list_keys_page(
|
||||||
|
&self,
|
||||||
|
query: &ProviderCatalogKeyListQuery,
|
||||||
|
) -> Result<StoredProviderCatalogKeyPage, DataLayerError> {
|
||||||
|
self.inner.list_keys_page(query).await
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn list_key_stats_by_provider_ids(
|
||||||
|
&self,
|
||||||
|
provider_ids: &[String],
|
||||||
|
) -> Result<Vec<StoredProviderCatalogKeyStats>, DataLayerError> {
|
||||||
|
let key = ProviderCatalogCacheKey::KeyStatsByProviderIds(normalize_ids(provider_ids));
|
||||||
|
match self
|
||||||
|
.get_or_load(key, || async move {
|
||||||
|
self.inner
|
||||||
|
.list_key_stats_by_provider_ids(provider_ids)
|
||||||
|
.await
|
||||||
|
.map(ProviderCatalogCacheValue::KeyStats)
|
||||||
|
})
|
||||||
|
.await?
|
||||||
|
{
|
||||||
|
ProviderCatalogCacheValue::KeyStats(items) => Ok(items),
|
||||||
|
_ => Ok(Vec::new()),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Clone, Debug, Eq, Hash, PartialEq)]
|
||||||
|
enum ProviderCatalogCacheKey {
|
||||||
|
Providers { active_only: bool },
|
||||||
|
ProvidersByIds(Vec<String>),
|
||||||
|
EndpointsByProviderIds(Vec<String>),
|
||||||
|
KeysByProviderIds(Vec<String>),
|
||||||
|
KeySummariesByProviderIds(Vec<String>),
|
||||||
|
KeyMaintenanceSummariesByProviderIds(Vec<String>),
|
||||||
|
KeyStatsByProviderIds(Vec<String>),
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Clone)]
|
||||||
|
enum ProviderCatalogCacheValue {
|
||||||
|
Providers(Vec<StoredProviderCatalogProvider>),
|
||||||
|
Endpoints(Vec<StoredProviderCatalogEndpoint>),
|
||||||
|
Keys(Vec<StoredProviderCatalogKey>),
|
||||||
|
KeyMaintenanceSummaries(Vec<StoredProviderCatalogKeyMaintenanceSummary>),
|
||||||
|
KeyStats(Vec<StoredProviderCatalogKeyStats>),
|
||||||
|
}
|
||||||
|
|
||||||
|
enum InflightRegistration {
|
||||||
|
Leader(u64),
|
||||||
|
Follower,
|
||||||
|
Bypass,
|
||||||
|
}
|
||||||
|
|
||||||
|
struct InflightGuard<'a> {
|
||||||
|
cache: &'a CachedProviderCatalogReadRepository,
|
||||||
|
key: Option<ProviderCatalogCacheKey>,
|
||||||
|
token: u64,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl<'a> InflightGuard<'a> {
|
||||||
|
fn new(
|
||||||
|
cache: &'a CachedProviderCatalogReadRepository,
|
||||||
|
key: ProviderCatalogCacheKey,
|
||||||
|
token: u64,
|
||||||
|
) -> Self {
|
||||||
|
Self {
|
||||||
|
cache,
|
||||||
|
key: Some(key),
|
||||||
|
token,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn finish(&mut self) {
|
||||||
|
if let Some(key) = self.key.take() {
|
||||||
|
self.cache.finish_inflight(&key, self.token);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Drop for InflightGuard<'_> {
|
||||||
|
fn drop(&mut self) {
|
||||||
|
self.finish();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn normalize_ids(ids: &[String]) -> Vec<String> {
|
||||||
|
let mut normalized = ids
|
||||||
|
.iter()
|
||||||
|
.map(|id| id.trim())
|
||||||
|
.filter(|id| !id.is_empty())
|
||||||
|
.map(ToOwned::to_owned)
|
||||||
|
.collect::<Vec<_>>();
|
||||||
|
normalized.sort();
|
||||||
|
normalized.dedup();
|
||||||
|
normalized
|
||||||
|
}
|
||||||
@@ -113,7 +113,7 @@ pub(crate) const ADMIN_MODULE_DEFINITIONS: &[AdminModuleDefinition] = &[
|
|||||||
AdminModuleDefinition {
|
AdminModuleDefinition {
|
||||||
name: "model_directives",
|
name: "model_directives",
|
||||||
display_name: "模型后缀参数",
|
display_name: "模型后缀参数",
|
||||||
description: "允许通过模型名后缀覆盖推理参数",
|
description: "允许通过模型名后缀覆盖推理参数或服务层级",
|
||||||
category: "integration",
|
category: "integration",
|
||||||
env_key: "MODEL_DIRECTIVES_AVAILABLE",
|
env_key: "MODEL_DIRECTIVES_AVAILABLE",
|
||||||
default_available: true,
|
default_available: true,
|
||||||
|
|||||||
@@ -2,7 +2,7 @@ use std::collections::{BTreeMap, BTreeSet};
|
|||||||
use std::time::{Duration, SystemTime, UNIX_EPOCH};
|
use std::time::{Duration, SystemTime, UNIX_EPOCH};
|
||||||
|
|
||||||
use aether_data_contracts::repository::pool_scores::{
|
use aether_data_contracts::repository::pool_scores::{
|
||||||
GetPoolMemberScoresByIdsQuery, ListPoolMemberProbeCandidatesQuery, PoolMemberHardState,
|
ListPoolMemberProbeCandidatesQuery, ListPoolMemberScoresQuery, PoolMemberHardState,
|
||||||
PoolMemberIdentity, PoolMemberProbeAttempt, PoolMemberProbeResult, PoolMemberProbeStatus,
|
PoolMemberIdentity, PoolMemberProbeAttempt, PoolMemberProbeResult, PoolMemberProbeStatus,
|
||||||
StoredPoolMemberScore, POOL_KIND_PROVIDER_KEY_POOL,
|
StoredPoolMemberScore, POOL_KIND_PROVIDER_KEY_POOL,
|
||||||
};
|
};
|
||||||
@@ -22,7 +22,6 @@ use crate::admin_api::{
|
|||||||
};
|
};
|
||||||
use crate::{AppState, GatewayError};
|
use crate::{AppState, GatewayError};
|
||||||
|
|
||||||
use crate::ai_serving::provider_key_pool_score_id;
|
|
||||||
use crate::ai_serving::provider_key_pool_score_scope;
|
use crate::ai_serving::provider_key_pool_score_scope;
|
||||||
use crate::handlers::shared::provider_pool::{
|
use crate::handlers::shared::provider_pool::{
|
||||||
admin_provider_pool_quota_probe_active_members_key, AdminProviderPoolConfig,
|
admin_provider_pool_quota_probe_active_members_key, AdminProviderPoolConfig,
|
||||||
@@ -38,6 +37,7 @@ const POOL_QUOTA_PROBE_REDIS_PREFIX: &str = "ap:quota_probe:last";
|
|||||||
const POOL_QUOTA_PROBE_DEFAULT_SCAN_INTERVAL_SECONDS: u64 = 60;
|
const POOL_QUOTA_PROBE_DEFAULT_SCAN_INTERVAL_SECONDS: u64 = 60;
|
||||||
const POOL_QUOTA_PROBE_MIN_SCAN_INTERVAL_SECONDS: u64 = 15;
|
const POOL_QUOTA_PROBE_MIN_SCAN_INTERVAL_SECONDS: u64 = 15;
|
||||||
const POOL_QUOTA_PROBE_DEFAULT_MAX_KEYS_PER_PROVIDER: usize = 50;
|
const POOL_QUOTA_PROBE_DEFAULT_MAX_KEYS_PER_PROVIDER: usize = 50;
|
||||||
|
const POOL_QUOTA_PROBE_PROVIDER_SCORE_READ_LIMIT: usize = 100_000;
|
||||||
const POOL_QUOTA_PROBE_DEFAULT_GLOBAL_CONCURRENCY: usize = 16;
|
const POOL_QUOTA_PROBE_DEFAULT_GLOBAL_CONCURRENCY: usize = 16;
|
||||||
const POOL_QUOTA_PROBE_PROVIDER_LOCK_TTL_MS: u64 = 30_000;
|
const POOL_QUOTA_PROBE_PROVIDER_LOCK_TTL_MS: u64 = 30_000;
|
||||||
const POOL_QUOTA_PROBE_BURST_TRIGGER_LOCK_TTL_MS: u64 = 30_000;
|
const POOL_QUOTA_PROBE_BURST_TRIGGER_LOCK_TTL_MS: u64 = 30_000;
|
||||||
@@ -386,21 +386,25 @@ async fn load_provider_key_account_scores(
|
|||||||
return BTreeMap::new();
|
return BTreeMap::new();
|
||||||
}
|
}
|
||||||
let scope = provider_key_pool_score_scope();
|
let scope = provider_key_pool_score_scope();
|
||||||
let score_ids = key_ids
|
let key_ids = key_ids.iter().map(String::as_str).collect::<BTreeSet<_>>();
|
||||||
.iter()
|
|
||||||
.map(|key_id| {
|
|
||||||
let identity =
|
|
||||||
PoolMemberIdentity::provider_api_key(provider_id.to_string(), key_id.clone());
|
|
||||||
provider_key_pool_score_id(&identity, &scope)
|
|
||||||
})
|
|
||||||
.collect::<Vec<_>>();
|
|
||||||
match state
|
match state
|
||||||
.data
|
.data
|
||||||
.get_pool_member_scores_by_ids(&GetPoolMemberScoresByIdsQuery { ids: score_ids })
|
.list_pool_member_scores(&ListPoolMemberScoresQuery {
|
||||||
|
pool_kind: POOL_KIND_PROVIDER_KEY_POOL.to_string(),
|
||||||
|
pool_id: provider_id.to_string(),
|
||||||
|
capability: Some(scope.capability),
|
||||||
|
scope_kind: Some(scope.scope_kind),
|
||||||
|
scope_id: scope.scope_id,
|
||||||
|
hard_states: Vec::new(),
|
||||||
|
probe_statuses: None,
|
||||||
|
offset: 0,
|
||||||
|
limit: POOL_QUOTA_PROBE_PROVIDER_SCORE_READ_LIMIT,
|
||||||
|
})
|
||||||
.await
|
.await
|
||||||
{
|
{
|
||||||
Ok(scores) => scores
|
Ok(scores) => scores
|
||||||
.into_iter()
|
.into_iter()
|
||||||
|
.filter(|score| key_ids.contains(score.member_id.as_str()))
|
||||||
.map(|score| (score.member_id.clone(), score))
|
.map(|score| (score.member_id.clone(), score))
|
||||||
.collect(),
|
.collect(),
|
||||||
Err(err) => {
|
Err(err) => {
|
||||||
|
|||||||
@@ -645,6 +645,11 @@ async fn record_health_success_effect(
|
|||||||
.circuit_breaker_by_format
|
.circuit_breaker_by_format
|
||||||
.as_ref()
|
.as_ref()
|
||||||
.and_then(|current| project_local_key_circuit_closed(Some(current), api_format));
|
.and_then(|current| project_local_key_circuit_closed(Some(current), api_format));
|
||||||
|
if current_key.health_by_format.as_ref() == Some(&health_by_format)
|
||||||
|
&& circuit_breaker_by_format.as_ref() == current_key.circuit_breaker_by_format.as_ref()
|
||||||
|
{
|
||||||
|
return;
|
||||||
|
}
|
||||||
let circuit_breaker_update = circuit_breaker_by_format
|
let circuit_breaker_update = circuit_breaker_by_format
|
||||||
.as_ref()
|
.as_ref()
|
||||||
.or(current_key.circuit_breaker_by_format.as_ref());
|
.or(current_key.circuit_breaker_by_format.as_ref());
|
||||||
|
|||||||
@@ -573,12 +573,14 @@ impl AppState {
|
|||||||
|
|
||||||
pub(crate) fn invalidate_provider_routing_caches(&self) {
|
pub(crate) fn invalidate_provider_routing_caches(&self) {
|
||||||
self.data.clear_minimal_candidate_selection_cache();
|
self.data.clear_minimal_candidate_selection_cache();
|
||||||
|
self.data.clear_provider_catalog_cache();
|
||||||
self.clear_provider_transport_snapshot_cache();
|
self.clear_provider_transport_snapshot_cache();
|
||||||
self.invalidate_scheduler_affinity_cache();
|
self.invalidate_scheduler_affinity_cache();
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(crate) fn invalidate_provider_health_routing_caches(&self) {
|
pub(crate) fn invalidate_provider_health_routing_caches(&self) {
|
||||||
self.data.clear_minimal_candidate_selection_cache();
|
self.data.clear_minimal_candidate_selection_cache();
|
||||||
|
self.data.clear_provider_catalog_cache();
|
||||||
self.clear_provider_transport_snapshot_cache();
|
self.clear_provider_transport_snapshot_cache();
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -88,15 +88,16 @@ pub(crate) async fn reasoning_model_directive_enabled_for_api_format_and_model(
|
|||||||
return false;
|
return false;
|
||||||
}
|
}
|
||||||
|
|
||||||
let Some(suffix) = requested_model.and_then(reasoning_suffix_from_model) else {
|
let Some(suffixes) = requested_model.and_then(model_directive_suffixes_from_model) else {
|
||||||
return false;
|
return false;
|
||||||
};
|
};
|
||||||
|
|
||||||
settings
|
let mappings = settings
|
||||||
.as_ref()
|
.as_ref()
|
||||||
.and_then(|settings| settings.api_format_mappings(&api_format))
|
.and_then(|settings| settings.api_format_mappings(&api_format));
|
||||||
.map(|mappings| mappings.contains_key(&suffix))
|
suffixes
|
||||||
.unwrap_or_else(|| DEFAULT_REASONING_SUFFIXES.contains(&suffix.as_str()))
|
.iter()
|
||||||
|
.all(|suffix| suffix_supported_for_api_format(&api_format, suffix, mappings.as_ref()))
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(crate) async fn reasoning_model_directive_mapping_for_api_format_and_model(
|
pub(crate) async fn reasoning_model_directive_mapping_for_api_format_and_model(
|
||||||
@@ -113,17 +114,17 @@ pub(crate) async fn reasoning_model_directive_mapping_for_api_format_and_model(
|
|||||||
{
|
{
|
||||||
return None;
|
return None;
|
||||||
}
|
}
|
||||||
let suffix = requested_model.and_then(reasoning_suffix_from_model)?;
|
let suffixes = requested_model.and_then(model_directive_suffixes_from_model)?;
|
||||||
let api_format = crate::ai_serving::normalize_api_format_alias(api_format);
|
let api_format = crate::ai_serving::normalize_api_format_alias(api_format);
|
||||||
let settings = read_reasoning_model_directive_settings(state).await;
|
let settings = read_reasoning_model_directive_settings(state).await;
|
||||||
settings
|
let mappings = settings
|
||||||
.as_ref()
|
.as_ref()
|
||||||
.and_then(|settings| settings.api_format_mappings(&api_format))
|
.and_then(|settings| settings.api_format_mappings(&api_format));
|
||||||
.and_then(|mappings| mappings.get(&suffix).cloned())
|
model_directive_mapping_for_suffixes(&api_format, &suffixes, mappings.as_ref())
|
||||||
.or_else(|| default_reasoning_mapping(&api_format, &suffix))
|
|
||||||
}
|
}
|
||||||
|
|
||||||
const DEFAULT_REASONING_SUFFIXES: &[&str] = &["low", "medium", "high", "xhigh", "max"];
|
const DEFAULT_MODEL_DIRECTIVE_SUFFIXES: &[&str] =
|
||||||
|
&["low", "medium", "high", "xhigh", "max", "fast"];
|
||||||
|
|
||||||
#[derive(Debug, Clone, Default)]
|
#[derive(Debug, Clone, Default)]
|
||||||
struct ReasoningModelDirectiveSettings {
|
struct ReasoningModelDirectiveSettings {
|
||||||
@@ -183,22 +184,66 @@ impl ReasoningModelDirectiveSettings {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
fn reasoning_suffix_from_model(model: &str) -> Option<String> {
|
fn model_directive_suffixes_from_model(model: &str) -> Option<Vec<String>> {
|
||||||
let model = model.trim();
|
let mut base_model = model.trim();
|
||||||
let (base_model, suffix) = model.rsplit_once('-')?;
|
let mut suffixes = Vec::new();
|
||||||
if base_model.trim().is_empty() {
|
let mut has_reasoning_effort = false;
|
||||||
|
let mut has_service_tier = false;
|
||||||
|
while let Some((candidate_base, suffix)) = base_model.rsplit_once('-') {
|
||||||
|
let Some(suffix) = normalize_reasoning_suffix(suffix) else {
|
||||||
|
break;
|
||||||
|
};
|
||||||
|
match model_directive_suffix_kind(&suffix)? {
|
||||||
|
ModelDirectiveSuffixKind::ReasoningEffort => {
|
||||||
|
if has_reasoning_effort {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
has_reasoning_effort = true;
|
||||||
|
}
|
||||||
|
ModelDirectiveSuffixKind::ServiceTier => {
|
||||||
|
if has_service_tier {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
has_service_tier = true;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
suffixes.push(suffix);
|
||||||
|
base_model = candidate_base.trim();
|
||||||
|
}
|
||||||
|
if base_model.is_empty() || suffixes.is_empty() {
|
||||||
return None;
|
return None;
|
||||||
}
|
}
|
||||||
normalize_reasoning_suffix(suffix)
|
suffixes.sort_by_key(|suffix| {
|
||||||
|
DEFAULT_MODEL_DIRECTIVE_SUFFIXES
|
||||||
|
.iter()
|
||||||
|
.position(|value| value == suffix)
|
||||||
|
.unwrap_or(usize::MAX)
|
||||||
|
});
|
||||||
|
Some(suffixes)
|
||||||
}
|
}
|
||||||
|
|
||||||
fn normalize_reasoning_suffix(suffix: &str) -> Option<String> {
|
fn normalize_reasoning_suffix(suffix: &str) -> Option<String> {
|
||||||
let normalized = suffix.trim().to_ascii_lowercase();
|
let normalized = suffix.trim().to_ascii_lowercase();
|
||||||
DEFAULT_REASONING_SUFFIXES
|
DEFAULT_MODEL_DIRECTIVE_SUFFIXES
|
||||||
.contains(&normalized.as_str())
|
.contains(&normalized.as_str())
|
||||||
.then_some(normalized)
|
.then_some(normalized)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
enum ModelDirectiveSuffixKind {
|
||||||
|
ReasoningEffort,
|
||||||
|
ServiceTier,
|
||||||
|
}
|
||||||
|
|
||||||
|
fn model_directive_suffix_kind(suffix: &str) -> Option<ModelDirectiveSuffixKind> {
|
||||||
|
match suffix {
|
||||||
|
"low" | "medium" | "high" | "xhigh" | "max" => {
|
||||||
|
Some(ModelDirectiveSuffixKind::ReasoningEffort)
|
||||||
|
}
|
||||||
|
"fast" => Some(ModelDirectiveSuffixKind::ServiceTier),
|
||||||
|
_ => None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
fn normalize_reasoning_mappings(
|
fn normalize_reasoning_mappings(
|
||||||
mappings: &serde_json::Map<String, serde_json::Value>,
|
mappings: &serde_json::Map<String, serde_json::Value>,
|
||||||
) -> serde_json::Map<String, serde_json::Value> {
|
) -> serde_json::Map<String, serde_json::Value> {
|
||||||
@@ -212,10 +257,16 @@ fn normalize_reasoning_mappings(
|
|||||||
|
|
||||||
fn default_reasoning_mapping(api_format: &str, suffix: &str) -> Option<serde_json::Value> {
|
fn default_reasoning_mapping(api_format: &str, suffix: &str) -> Option<serde_json::Value> {
|
||||||
match api_format {
|
match api_format {
|
||||||
|
"openai:chat" if suffix == "fast" => {
|
||||||
|
Some(serde_json::json!({ "service_tier": "priority" }))
|
||||||
|
}
|
||||||
"openai:chat" => {
|
"openai:chat" => {
|
||||||
let effort = openai_reasoning_effort_value(suffix)?;
|
let effort = openai_reasoning_effort_value(suffix)?;
|
||||||
Some(serde_json::json!({ "reasoning_effort": effort }))
|
Some(serde_json::json!({ "reasoning_effort": effort }))
|
||||||
}
|
}
|
||||||
|
"openai:responses" | "openai:responses:compact" if suffix == "fast" => {
|
||||||
|
Some(serde_json::json!({ "service_tier": "priority" }))
|
||||||
|
}
|
||||||
"openai:responses" | "openai:responses:compact" => {
|
"openai:responses" | "openai:responses:compact" => {
|
||||||
let effort = openai_reasoning_effort_value(suffix)?;
|
let effort = openai_reasoning_effort_value(suffix)?;
|
||||||
Some(serde_json::json!({ "reasoning": { "effort": effort } }))
|
Some(serde_json::json!({ "reasoning": { "effort": effort } }))
|
||||||
@@ -251,6 +302,54 @@ fn default_reasoning_mapping(api_format: &str, suffix: &str) -> Option<serde_jso
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn suffix_supported_for_api_format(
|
||||||
|
api_format: &str,
|
||||||
|
suffix: &str,
|
||||||
|
mappings: Option<&serde_json::Map<String, serde_json::Value>>,
|
||||||
|
) -> bool {
|
||||||
|
mappings
|
||||||
|
.map(|mappings| mappings.contains_key(suffix))
|
||||||
|
.unwrap_or_else(|| default_reasoning_mapping(api_format, suffix).is_some())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn model_directive_mapping_for_suffixes(
|
||||||
|
api_format: &str,
|
||||||
|
suffixes: &[String],
|
||||||
|
mappings: Option<&serde_json::Map<String, serde_json::Value>>,
|
||||||
|
) -> Option<serde_json::Value> {
|
||||||
|
let mut combined = serde_json::json!({});
|
||||||
|
for suffix in suffixes {
|
||||||
|
let mapping = mappings
|
||||||
|
.and_then(|mappings| mappings.get(suffix).cloned())
|
||||||
|
.or_else(|| {
|
||||||
|
mappings
|
||||||
|
.is_none()
|
||||||
|
.then(|| default_reasoning_mapping(api_format, suffix))
|
||||||
|
.flatten()
|
||||||
|
})?;
|
||||||
|
deep_merge_json(&mut combined, &mapping);
|
||||||
|
}
|
||||||
|
Some(combined)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn deep_merge_json(target: &mut serde_json::Value, patch: &serde_json::Value) {
|
||||||
|
match (target, patch) {
|
||||||
|
(serde_json::Value::Object(target_object), serde_json::Value::Object(patch_object)) => {
|
||||||
|
for (key, patch_value) in patch_object {
|
||||||
|
match target_object.get_mut(key) {
|
||||||
|
Some(target_value) => deep_merge_json(target_value, patch_value),
|
||||||
|
None => {
|
||||||
|
target_object.insert(key.clone(), patch_value.clone());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
(target, patch) => {
|
||||||
|
*target = patch.clone();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
fn openai_reasoning_effort_value(suffix: &str) -> Option<&'static str> {
|
fn openai_reasoning_effort_value(suffix: &str) -> Option<&'static str> {
|
||||||
match suffix {
|
match suffix {
|
||||||
"low" => Some("low"),
|
"low" => Some("low"),
|
||||||
@@ -297,7 +396,11 @@ fn parse_reasoning_model_directive_settings(
|
|||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use super::parse_reasoning_model_directive_settings;
|
use super::{
|
||||||
|
default_reasoning_mapping, model_directive_mapping_for_suffixes,
|
||||||
|
model_directive_suffixes_from_model, parse_reasoning_model_directive_settings,
|
||||||
|
suffix_supported_for_api_format,
|
||||||
|
};
|
||||||
use serde_json::json;
|
use serde_json::json;
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
@@ -332,4 +435,64 @@ mod tests {
|
|||||||
);
|
);
|
||||||
assert_eq!(settings.api_format_enabled("gemini:generate_content"), None);
|
assert_eq!(settings.api_format_enabled("gemini:generate_content"), None);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn default_fast_suffix_maps_to_openai_priority_service_tier() {
|
||||||
|
assert_eq!(
|
||||||
|
default_reasoning_mapping("openai:chat", "fast"),
|
||||||
|
Some(json!({ "service_tier": "priority" }))
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
default_reasoning_mapping("openai:responses", "fast"),
|
||||||
|
Some(json!({ "service_tier": "priority" }))
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
default_reasoning_mapping("openai:responses:compact", "fast"),
|
||||||
|
Some(json!({ "service_tier": "priority" }))
|
||||||
|
);
|
||||||
|
assert_eq!(default_reasoning_mapping("claude:messages", "fast"), None);
|
||||||
|
assert_eq!(
|
||||||
|
default_reasoning_mapping("gemini:generate_content", "fast"),
|
||||||
|
None
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn combined_suffixes_are_order_insensitive() {
|
||||||
|
let expected = Some(vec!["xhigh".to_string(), "fast".to_string()]);
|
||||||
|
assert_eq!(
|
||||||
|
model_directive_suffixes_from_model("gpt-5.4-fast-xhigh"),
|
||||||
|
expected
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
model_directive_suffixes_from_model("gpt-5.4-xhigh-fast"),
|
||||||
|
expected
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
model_directive_mapping_for_suffixes(
|
||||||
|
"openai:chat",
|
||||||
|
expected.as_ref().expect("suffixes should parse"),
|
||||||
|
None,
|
||||||
|
),
|
||||||
|
Some(json!({
|
||||||
|
"reasoning_effort": "xhigh",
|
||||||
|
"service_tier": "priority"
|
||||||
|
}))
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn combined_suffix_requires_each_suffix_to_support_api_format() {
|
||||||
|
let suffixes = model_directive_suffixes_from_model("gpt-5.4-xhigh-fast").expect("suffixes");
|
||||||
|
assert!(suffixes
|
||||||
|
.iter()
|
||||||
|
.all(|suffix| { suffix_supported_for_api_format("openai:responses", suffix, None) }));
|
||||||
|
assert!(!suffixes
|
||||||
|
.iter()
|
||||||
|
.all(|suffix| { suffix_supported_for_api_format("claude:messages", suffix, None) }));
|
||||||
|
assert_eq!(
|
||||||
|
model_directive_mapping_for_suffixes("claude:messages", &suffixes, None),
|
||||||
|
None
|
||||||
|
);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1729,7 +1729,8 @@ pub fn admin_system_config_default_value(key: &str) -> Option<serde_json::Value>
|
|||||||
"medium": { "reasoning_effort": "medium" },
|
"medium": { "reasoning_effort": "medium" },
|
||||||
"high": { "reasoning_effort": "high" },
|
"high": { "reasoning_effort": "high" },
|
||||||
"xhigh": { "reasoning_effort": "xhigh" },
|
"xhigh": { "reasoning_effort": "xhigh" },
|
||||||
"max": { "reasoning_effort": "xhigh" }
|
"max": { "reasoning_effort": "xhigh" },
|
||||||
|
"fast": { "service_tier": "priority" }
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
"openai:responses": {
|
"openai:responses": {
|
||||||
@@ -1739,7 +1740,8 @@ pub fn admin_system_config_default_value(key: &str) -> Option<serde_json::Value>
|
|||||||
"medium": { "reasoning": { "effort": "medium" } },
|
"medium": { "reasoning": { "effort": "medium" } },
|
||||||
"high": { "reasoning": { "effort": "high" } },
|
"high": { "reasoning": { "effort": "high" } },
|
||||||
"xhigh": { "reasoning": { "effort": "xhigh" } },
|
"xhigh": { "reasoning": { "effort": "xhigh" } },
|
||||||
"max": { "reasoning": { "effort": "xhigh" } }
|
"max": { "reasoning": { "effort": "xhigh" } },
|
||||||
|
"fast": { "service_tier": "priority" }
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
"openai:responses:compact": {
|
"openai:responses:compact": {
|
||||||
@@ -1749,7 +1751,8 @@ pub fn admin_system_config_default_value(key: &str) -> Option<serde_json::Value>
|
|||||||
"medium": { "reasoning": { "effort": "medium" } },
|
"medium": { "reasoning": { "effort": "medium" } },
|
||||||
"high": { "reasoning": { "effort": "high" } },
|
"high": { "reasoning": { "effort": "high" } },
|
||||||
"xhigh": { "reasoning": { "effort": "xhigh" } },
|
"xhigh": { "reasoning": { "effort": "xhigh" } },
|
||||||
"max": { "reasoning": { "effort": "xhigh" } }
|
"max": { "reasoning": { "effort": "xhigh" } },
|
||||||
|
"fast": { "service_tier": "priority" }
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
"claude:messages": {
|
"claude:messages": {
|
||||||
|
|||||||
@@ -77,7 +77,7 @@ pub use crate::formats::shared::model_directives::{
|
|||||||
apply_model_directive_overrides_from_request, claude_model_uses_adaptive_effort,
|
apply_model_directive_overrides_from_request, claude_model_uses_adaptive_effort,
|
||||||
extract_gemini_model_from_path, gemini_model_uses_thinking_level, model_directive_base_model,
|
extract_gemini_model_from_path, gemini_model_uses_thinking_level, model_directive_base_model,
|
||||||
normalize_model_directive_model, parse_model_directive, ModelDirective, ModelOverride,
|
normalize_model_directive_model, parse_model_directive, ModelDirective, ModelOverride,
|
||||||
ReasoningEffort,
|
ReasoningEffort, ServiceTier,
|
||||||
};
|
};
|
||||||
pub use crate::formats::shared::passthrough::{
|
pub use crate::formats::shared::passthrough::{
|
||||||
resolve_stream_spec as resolve_local_same_format_stream_spec,
|
resolve_stream_spec as resolve_local_same_format_stream_spec,
|
||||||
|
|||||||
@@ -9,6 +9,7 @@ pub struct ModelDirective {
|
|||||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||||
pub enum ModelOverride {
|
pub enum ModelOverride {
|
||||||
ReasoningEffort(ReasoningEffort),
|
ReasoningEffort(ReasoningEffort),
|
||||||
|
ServiceTier(ServiceTier),
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||||
@@ -79,19 +80,90 @@ impl ReasoningEffort {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||||
|
pub enum ServiceTier {
|
||||||
|
Priority,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ServiceTier {
|
||||||
|
pub fn parse(value: &str) -> Option<Self> {
|
||||||
|
match value.trim().to_ascii_lowercase().as_str() {
|
||||||
|
"fast" => Some(Self::Priority),
|
||||||
|
_ => None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn as_openai_value(self) -> &'static str {
|
||||||
|
match self {
|
||||||
|
Self::Priority => "priority",
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
pub fn parse_model_directive(model: &str) -> Option<ModelDirective> {
|
pub fn parse_model_directive(model: &str) -> Option<ModelDirective> {
|
||||||
let model = model.trim();
|
let (base_model, overrides) = parse_model_directive_parts(model)?;
|
||||||
let (base_model, suffix) = model.rsplit_once('-')?;
|
Some(ModelDirective {
|
||||||
let base_model = base_model.trim();
|
base_model,
|
||||||
|
overrides,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn parse_model_directive_parts(model: &str) -> Option<(String, Vec<ModelOverride>)> {
|
||||||
|
let mut base_model = model.trim();
|
||||||
|
let mut overrides = ModelOverrideAccumulator::default();
|
||||||
|
while let Some((candidate_base, suffix)) = base_model.rsplit_once('-') {
|
||||||
|
let Some(override_item) = parse_model_override(suffix) else {
|
||||||
|
break;
|
||||||
|
};
|
||||||
|
overrides.insert(override_item)?;
|
||||||
|
base_model = candidate_base.trim();
|
||||||
|
}
|
||||||
if base_model.is_empty() {
|
if base_model.is_empty() {
|
||||||
return None;
|
return None;
|
||||||
}
|
}
|
||||||
|
let overrides = overrides.into_overrides()?;
|
||||||
|
Some((base_model.to_string(), overrides))
|
||||||
|
}
|
||||||
|
|
||||||
let reasoning_effort = ReasoningEffort::parse(suffix)?;
|
fn parse_model_override(suffix: &str) -> Option<ModelOverride> {
|
||||||
Some(ModelDirective {
|
ReasoningEffort::parse(suffix)
|
||||||
base_model: base_model.to_string(),
|
.map(ModelOverride::ReasoningEffort)
|
||||||
overrides: vec![ModelOverride::ReasoningEffort(reasoning_effort)],
|
.or_else(|| ServiceTier::parse(suffix).map(ModelOverride::ServiceTier))
|
||||||
})
|
}
|
||||||
|
|
||||||
|
#[derive(Default)]
|
||||||
|
struct ModelOverrideAccumulator {
|
||||||
|
reasoning_effort: Option<ReasoningEffort>,
|
||||||
|
service_tier: Option<ServiceTier>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ModelOverrideAccumulator {
|
||||||
|
fn insert(&mut self, override_item: ModelOverride) -> Option<()> {
|
||||||
|
match override_item {
|
||||||
|
ModelOverride::ReasoningEffort(value) => {
|
||||||
|
if self.reasoning_effort.replace(value).is_some() {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
ModelOverride::ServiceTier(value) => {
|
||||||
|
if self.service_tier.replace(value).is_some() {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Some(())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn into_overrides(self) -> Option<Vec<ModelOverride>> {
|
||||||
|
let mut overrides = Vec::new();
|
||||||
|
if let Some(reasoning_effort) = self.reasoning_effort {
|
||||||
|
overrides.push(ModelOverride::ReasoningEffort(reasoning_effort));
|
||||||
|
}
|
||||||
|
if let Some(service_tier) = self.service_tier {
|
||||||
|
overrides.push(ModelOverride::ServiceTier(service_tier));
|
||||||
|
}
|
||||||
|
(!overrides.is_empty()).then_some(overrides)
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn model_directive_base_model(model: &str) -> Option<String> {
|
pub fn model_directive_base_model(model: &str) -> Option<String> {
|
||||||
@@ -149,18 +221,23 @@ pub fn apply_model_directive_overrides_from_model(
|
|||||||
source_model: &str,
|
source_model: &str,
|
||||||
) -> Option<ModelDirective> {
|
) -> Option<ModelDirective> {
|
||||||
let directive = parse_model_directive(source_model)?;
|
let directive = parse_model_directive(source_model)?;
|
||||||
|
let mut patched_body = provider_request_body.clone();
|
||||||
for override_item in &directive.overrides {
|
for override_item in &directive.overrides {
|
||||||
match override_item {
|
match override_item {
|
||||||
ModelOverride::ReasoningEffort(effort) => {
|
ModelOverride::ReasoningEffort(effort) => {
|
||||||
apply_reasoning_effort_override(
|
apply_reasoning_effort_override(
|
||||||
provider_request_body,
|
&mut patched_body,
|
||||||
provider_api_format,
|
provider_api_format,
|
||||||
provider_model,
|
provider_model,
|
||||||
*effort,
|
*effort,
|
||||||
)?;
|
)?;
|
||||||
}
|
}
|
||||||
|
ModelOverride::ServiceTier(tier) => {
|
||||||
|
apply_service_tier_override(&mut patched_body, provider_api_format, *tier)?;
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
*provider_request_body = patched_body;
|
||||||
Some(directive)
|
Some(directive)
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -215,6 +292,21 @@ fn apply_reasoning_effort_override(
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn apply_service_tier_override(
|
||||||
|
provider_request_body: &mut Value,
|
||||||
|
provider_api_format: &str,
|
||||||
|
tier: ServiceTier,
|
||||||
|
) -> Option<()> {
|
||||||
|
match crate::normalize_api_format_alias(provider_api_format).as_str() {
|
||||||
|
"openai:chat" | "openai:responses" | "openai:responses:compact" => set_object_string(
|
||||||
|
provider_request_body,
|
||||||
|
"service_tier",
|
||||||
|
tier.as_openai_value(),
|
||||||
|
),
|
||||||
|
_ => None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
fn set_object_string(body: &mut Value, key: &str, value: &str) -> Option<()> {
|
fn set_object_string(body: &mut Value, key: &str, value: &str) -> Option<()> {
|
||||||
body.as_object_mut()?
|
body.as_object_mut()?
|
||||||
.insert(key.to_string(), Value::String(value.to_string()));
|
.insert(key.to_string(), Value::String(value.to_string()));
|
||||||
@@ -368,7 +460,7 @@ mod tests {
|
|||||||
|
|
||||||
use super::{
|
use super::{
|
||||||
apply_model_directive_overrides_from_model, parse_model_directive, ModelDirective,
|
apply_model_directive_overrides_from_model, parse_model_directive, ModelDirective,
|
||||||
ModelOverride, ReasoningEffort,
|
ModelOverride, ReasoningEffort, ServiceTier,
|
||||||
};
|
};
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
@@ -389,12 +481,37 @@ mod tests {
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn parses_supported_service_tier_suffixes() {
|
||||||
|
assert_eq!(
|
||||||
|
parse_model_directive("gpt-5.4-fast"),
|
||||||
|
Some(ModelDirective {
|
||||||
|
base_model: "gpt-5.4".to_string(),
|
||||||
|
overrides: vec![ModelOverride::ServiceTier(ServiceTier::Priority)],
|
||||||
|
})
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn parses_combined_suffixes_in_canonical_order() {
|
||||||
|
let expected = Some(ModelDirective {
|
||||||
|
base_model: "gpt-5.4".to_string(),
|
||||||
|
overrides: vec![
|
||||||
|
ModelOverride::ReasoningEffort(ReasoningEffort::XHigh),
|
||||||
|
ModelOverride::ServiceTier(ServiceTier::Priority),
|
||||||
|
],
|
||||||
|
});
|
||||||
|
assert_eq!(parse_model_directive("gpt-5.4-fast-xhigh"), expected);
|
||||||
|
assert_eq!(parse_model_directive("gpt-5.4-xhigh-fast"), expected);
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn ignores_unknown_or_incomplete_suffixes() {
|
fn ignores_unknown_or_incomplete_suffixes() {
|
||||||
assert_eq!(parse_model_directive("gpt-5.4-ultra"), None);
|
assert_eq!(parse_model_directive("gpt-5.4-ultra"), None);
|
||||||
assert_eq!(parse_model_directive("gpt-5.4"), None);
|
assert_eq!(parse_model_directive("gpt-5.4"), None);
|
||||||
assert_eq!(parse_model_directive("-high"), None);
|
assert_eq!(parse_model_directive("-high"), None);
|
||||||
assert_eq!(parse_model_directive("gpt-5.4-high-json"), None);
|
assert_eq!(parse_model_directive("gpt-5.4-high-json"), None);
|
||||||
|
assert_eq!(parse_model_directive("gpt-5.4-low-high"), None);
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
@@ -446,4 +563,65 @@ mod tests {
|
|||||||
2048
|
2048
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn applies_fast_suffix_to_openai_service_tier() {
|
||||||
|
let mut openai_chat = json!({"model": "gpt-5-upstream"});
|
||||||
|
apply_model_directive_overrides_from_model(
|
||||||
|
&mut openai_chat,
|
||||||
|
"openai:chat",
|
||||||
|
"gpt-5-upstream",
|
||||||
|
"gpt-5.4-fast",
|
||||||
|
)
|
||||||
|
.expect("directive should apply");
|
||||||
|
assert_eq!(openai_chat["service_tier"], "priority");
|
||||||
|
|
||||||
|
let mut responses = json!({"model": "gpt-5-upstream"});
|
||||||
|
apply_model_directive_overrides_from_model(
|
||||||
|
&mut responses,
|
||||||
|
"openai:responses",
|
||||||
|
"gpt-5-upstream",
|
||||||
|
"gpt-5.4-fast",
|
||||||
|
)
|
||||||
|
.expect("directive should apply");
|
||||||
|
assert_eq!(responses["service_tier"], "priority");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn applies_combined_suffixes_to_openai_body() {
|
||||||
|
let mut openai_chat = json!({"model": "gpt-5-upstream", "reasoning_effort": "low"});
|
||||||
|
apply_model_directive_overrides_from_model(
|
||||||
|
&mut openai_chat,
|
||||||
|
"openai:chat",
|
||||||
|
"gpt-5-upstream",
|
||||||
|
"gpt-5.4-fast-xhigh",
|
||||||
|
)
|
||||||
|
.expect("directive should apply");
|
||||||
|
assert_eq!(openai_chat["reasoning_effort"], "xhigh");
|
||||||
|
assert_eq!(openai_chat["service_tier"], "priority");
|
||||||
|
|
||||||
|
let mut reversed = json!({"model": "gpt-5-upstream", "reasoning_effort": "low"});
|
||||||
|
apply_model_directive_overrides_from_model(
|
||||||
|
&mut reversed,
|
||||||
|
"openai:chat",
|
||||||
|
"gpt-5-upstream",
|
||||||
|
"gpt-5.4-xhigh-fast",
|
||||||
|
)
|
||||||
|
.expect("directive should apply");
|
||||||
|
assert_eq!(reversed, openai_chat);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn unsupported_combined_suffix_leaves_body_unchanged() {
|
||||||
|
let mut claude = json!({"model": "claude-sonnet-4-5"});
|
||||||
|
let original = claude.clone();
|
||||||
|
assert!(apply_model_directive_overrides_from_model(
|
||||||
|
&mut claude,
|
||||||
|
"claude:messages",
|
||||||
|
"claude-sonnet-4-5",
|
||||||
|
"gpt-5.4-fast-xhigh",
|
||||||
|
)
|
||||||
|
.is_none());
|
||||||
|
assert_eq!(claude, original);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -25,7 +25,7 @@ pub use formats::shared::model_directives::{
|
|||||||
apply_model_directive_overrides_from_request, claude_model_uses_adaptive_effort,
|
apply_model_directive_overrides_from_request, claude_model_uses_adaptive_effort,
|
||||||
extract_gemini_model_from_path, gemini_model_uses_thinking_level, model_directive_base_model,
|
extract_gemini_model_from_path, gemini_model_uses_thinking_level, model_directive_base_model,
|
||||||
normalize_model_directive_model, parse_model_directive, ModelDirective, ModelOverride,
|
normalize_model_directive_model, parse_model_directive, ModelDirective, ModelOverride,
|
||||||
ReasoningEffort,
|
ReasoningEffort, ServiceTier,
|
||||||
};
|
};
|
||||||
pub use formats::shared::request::{
|
pub use formats::shared::request::{
|
||||||
endpoint_config_forces_upstream_stream_policy, enforce_request_body_stream_field,
|
endpoint_config_forces_upstream_stream_policy, enforce_request_body_stream_field,
|
||||||
|
|||||||
@@ -580,6 +580,8 @@ impl StoredProviderCatalogKeyStats {
|
|||||||
|
|
||||||
#[async_trait]
|
#[async_trait]
|
||||||
pub trait ProviderCatalogReadRepository: Send + Sync {
|
pub trait ProviderCatalogReadRepository: Send + Sync {
|
||||||
|
fn clear_local_cache(&self) {}
|
||||||
|
|
||||||
async fn list_providers(
|
async fn list_providers(
|
||||||
&self,
|
&self,
|
||||||
active_only: bool,
|
active_only: bool,
|
||||||
|
|||||||
@@ -418,19 +418,33 @@ impl StoredRequestUsageAudit {
|
|||||||
}
|
}
|
||||||
|
|
||||||
pub fn provider_reasoning_effort(&self) -> Option<String> {
|
pub fn provider_reasoning_effort(&self) -> Option<String> {
|
||||||
|
if self
|
||||||
|
.provider_request_body
|
||||||
|
.as_ref()
|
||||||
|
.and_then(Value::as_object)
|
||||||
|
.is_some()
|
||||||
|
{
|
||||||
|
return extract_provider_reasoning_effort_from_body(
|
||||||
|
self.provider_request_body.as_ref(),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
self.request_metadata_string(PROVIDER_REASONING_EFFORT_METADATA_KEY)
|
self.request_metadata_string(PROVIDER_REASONING_EFFORT_METADATA_KEY)
|
||||||
.and_then(normalize_provider_reasoning_effort)
|
.and_then(normalize_provider_reasoning_effort)
|
||||||
.or_else(|| {
|
|
||||||
extract_provider_reasoning_effort_from_body(self.provider_request_body.as_ref())
|
|
||||||
})
|
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn provider_service_tier(&self) -> Option<String> {
|
pub fn provider_service_tier(&self) -> Option<String> {
|
||||||
|
if self
|
||||||
|
.provider_request_body
|
||||||
|
.as_ref()
|
||||||
|
.and_then(Value::as_object)
|
||||||
|
.is_some()
|
||||||
|
{
|
||||||
|
return extract_provider_service_tier_from_body(self.provider_request_body.as_ref());
|
||||||
|
}
|
||||||
|
|
||||||
self.request_metadata_string(PROVIDER_SERVICE_TIER_METADATA_KEY)
|
self.request_metadata_string(PROVIDER_SERVICE_TIER_METADATA_KEY)
|
||||||
.and_then(normalize_provider_service_tier)
|
.and_then(normalize_provider_service_tier)
|
||||||
.or_else(|| {
|
|
||||||
extract_provider_service_tier_from_body(self.provider_request_body.as_ref())
|
|
||||||
})
|
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn body_ref(&self, field: UsageBodyField) -> Option<&str> {
|
pub fn body_ref(&self, field: UsageBodyField) -> Option<&str> {
|
||||||
@@ -2398,7 +2412,7 @@ mod tests {
|
|||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn provider_reasoning_effort_reads_metadata_and_provider_request_body() {
|
fn provider_reasoning_effort_prefers_provider_request_body_over_metadata() {
|
||||||
let mut usage = sample_usage();
|
let mut usage = sample_usage();
|
||||||
usage.provider_request_body = Some(json!({
|
usage.provider_request_body = Some(json!({
|
||||||
"reasoning": { "effort": "XHigh" },
|
"reasoning": { "effort": "XHigh" },
|
||||||
@@ -2413,6 +2427,10 @@ mod tests {
|
|||||||
"provider_service_tier": "standard"
|
"provider_service_tier": "standard"
|
||||||
}));
|
}));
|
||||||
|
|
||||||
|
assert_eq!(usage.provider_reasoning_effort().as_deref(), Some("xhigh"));
|
||||||
|
assert_eq!(usage.provider_service_tier().as_deref(), Some("priority"));
|
||||||
|
|
||||||
|
usage.provider_request_body = None;
|
||||||
assert_eq!(usage.provider_reasoning_effort().as_deref(), Some("max"));
|
assert_eq!(usage.provider_reasoning_effort().as_deref(), Some("max"));
|
||||||
assert_eq!(usage.provider_service_tier().as_deref(), Some("standard"));
|
assert_eq!(usage.provider_service_tier().as_deref(), Some("standard"));
|
||||||
|
|
||||||
@@ -2424,6 +2442,17 @@ mod tests {
|
|||||||
|
|
||||||
assert_eq!(usage.provider_reasoning_effort().as_deref(), Some("high"));
|
assert_eq!(usage.provider_reasoning_effort().as_deref(), Some("high"));
|
||||||
assert_eq!(usage.provider_service_tier().as_deref(), Some("priority"));
|
assert_eq!(usage.provider_service_tier().as_deref(), Some("priority"));
|
||||||
|
|
||||||
|
usage.request_metadata = Some(json!({
|
||||||
|
"provider_reasoning_effort": "max",
|
||||||
|
"provider_service_tier": "standard"
|
||||||
|
}));
|
||||||
|
usage.provider_request_body = Some(json!({
|
||||||
|
"model": "gpt-5"
|
||||||
|
}));
|
||||||
|
|
||||||
|
assert_eq!(usage.provider_reasoning_effort(), None);
|
||||||
|
assert_eq!(usage.provider_service_tier(), None);
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
|
|||||||
+15
@@ -0,0 +1,15 @@
|
|||||||
|
SET @aether_provider_key_name_index_sql := IF(
|
||||||
|
(
|
||||||
|
SELECT COUNT(*)
|
||||||
|
FROM information_schema.statistics
|
||||||
|
WHERE table_schema = DATABASE()
|
||||||
|
AND table_name = 'provider_api_keys'
|
||||||
|
AND index_name = 'idx_provider_api_keys_provider_name_id'
|
||||||
|
) = 0,
|
||||||
|
'CREATE INDEX idx_provider_api_keys_provider_name_id ON provider_api_keys (provider_id, name, id)',
|
||||||
|
'DO 0'
|
||||||
|
);
|
||||||
|
|
||||||
|
PREPARE aether_provider_key_name_index_stmt FROM @aether_provider_key_name_index_sql;
|
||||||
|
EXECUTE aether_provider_key_name_index_stmt;
|
||||||
|
DEALLOCATE PREPARE aether_provider_key_name_index_stmt;
|
||||||
+2
@@ -0,0 +1,2 @@
|
|||||||
|
CREATE INDEX IF NOT EXISTS idx_provider_api_keys_provider_name_id
|
||||||
|
ON public.provider_api_keys USING btree (provider_id, name, id);
|
||||||
+2
@@ -0,0 +1,2 @@
|
|||||||
|
CREATE INDEX IF NOT EXISTS idx_provider_api_keys_provider_name_id
|
||||||
|
ON provider_api_keys (provider_id, name, id);
|
||||||
@@ -197,6 +197,14 @@ CREATE INDEX IF NOT EXISTS idx_provider_api_keys_provider_default_sort ON public
|
|||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
--
|
||||||
|
-- Name: idx_provider_api_keys_provider_name_id; Type: INDEX; Schema: public; Owner: -
|
||||||
|
--
|
||||||
|
|
||||||
|
CREATE INDEX IF NOT EXISTS idx_provider_api_keys_provider_name_id ON public.provider_api_keys USING btree (provider_id, name, id);
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
--
|
--
|
||||||
-- Name: idx_provider_api_keys_provider_active_priority_id; Type: INDEX; Schema: public; Owner: -
|
-- Name: idx_provider_api_keys_provider_active_priority_id; Type: INDEX; Schema: public; Owner: -
|
||||||
--
|
--
|
||||||
|
|||||||
@@ -7,7 +7,7 @@ use tracing::info;
|
|||||||
// Generated by build.rs from schema/bootstrap/postgres.
|
// Generated by build.rs from schema/bootstrap/postgres.
|
||||||
pub(crate) static EMPTY_DATABASE_SNAPSHOT_SQL: &str =
|
pub(crate) static EMPTY_DATABASE_SNAPSHOT_SQL: &str =
|
||||||
include_str!(concat!(env!("OUT_DIR"), "/empty_database_snapshot.sql"));
|
include_str!(concat!(env!("OUT_DIR"), "/empty_database_snapshot.sql"));
|
||||||
pub(crate) const EMPTY_DATABASE_SNAPSHOT_CUTOFF_VERSION: i64 = 20260524000000;
|
pub(crate) const EMPTY_DATABASE_SNAPSHOT_CUTOFF_VERSION: i64 = 20260527000000;
|
||||||
|
|
||||||
const PUBLIC_BASE_TABLE_COUNT_SQL: &str = r#"
|
const PUBLIC_BASE_TABLE_COUNT_SQL: &str = r#"
|
||||||
SELECT COUNT(*)::BIGINT
|
SELECT COUNT(*)::BIGINT
|
||||||
|
|||||||
@@ -314,6 +314,7 @@ fn empty_database_snapshot_covers_current_cutoff_versions() {
|
|||||||
20260520010000,
|
20260520010000,
|
||||||
20260522000000,
|
20260522000000,
|
||||||
20260524000000,
|
20260524000000,
|
||||||
|
20260527000000,
|
||||||
]
|
]
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
@@ -390,6 +391,7 @@ fn empty_database_snapshot_sql_includes_usage_body_blobs_and_audit_admin_role()
|
|||||||
assert!(EMPTY_DATABASE_SNAPSHOT_SQL.contains("ix_usage_counter_deltas_unprocessed"));
|
assert!(EMPTY_DATABASE_SNAPSHOT_SQL.contains("ix_usage_counter_deltas_unprocessed"));
|
||||||
assert!(EMPTY_DATABASE_SNAPSHOT_SQL.contains("idx_entitlement_usage_entitlement_date"));
|
assert!(EMPTY_DATABASE_SNAPSHOT_SQL.contains("idx_entitlement_usage_entitlement_date"));
|
||||||
assert!(EMPTY_DATABASE_SNAPSHOT_SQL.contains("idx_provider_api_keys_provider_default_sort"));
|
assert!(EMPTY_DATABASE_SNAPSHOT_SQL.contains("idx_provider_api_keys_provider_default_sort"));
|
||||||
|
assert!(EMPTY_DATABASE_SNAPSHOT_SQL.contains("idx_provider_api_keys_provider_name_id"));
|
||||||
assert!(
|
assert!(
|
||||||
EMPTY_DATABASE_SNAPSHOT_SQL.contains("idx_provider_api_keys_provider_active_priority_id")
|
EMPTY_DATABASE_SNAPSHOT_SQL.contains("idx_provider_api_keys_provider_active_priority_id")
|
||||||
);
|
);
|
||||||
@@ -668,6 +670,7 @@ fn mysql_and_sqlite_migrations_include_enabled_incrementals() {
|
|||||||
20260520000000,
|
20260520000000,
|
||||||
20260520010000,
|
20260520010000,
|
||||||
20260524000000,
|
20260524000000,
|
||||||
|
20260527000000,
|
||||||
]
|
]
|
||||||
);
|
);
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
@@ -692,6 +695,7 @@ fn mysql_and_sqlite_migrations_include_enabled_incrementals() {
|
|||||||
20260520000000,
|
20260520000000,
|
||||||
20260520010000,
|
20260520010000,
|
||||||
20260524000000,
|
20260524000000,
|
||||||
|
20260527000000,
|
||||||
]
|
]
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
@@ -1217,6 +1221,7 @@ fn pending_migrations_from_applied_skips_versions_already_applied() {
|
|||||||
20260520010000,
|
20260520010000,
|
||||||
20260522000000,
|
20260522000000,
|
||||||
20260524000000,
|
20260524000000,
|
||||||
|
20260527000000,
|
||||||
]
|
]
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -341,14 +341,13 @@ impl PoolScoreReadRepository for PostgresPoolMemberScoreRepository {
|
|||||||
if query.ids.is_empty() {
|
if query.ids.is_empty() {
|
||||||
return Ok(Vec::new());
|
return Ok(Vec::new());
|
||||||
}
|
}
|
||||||
let mut builder = QueryBuilder::<Postgres>::new(SCORE_COLUMNS);
|
let rows = sqlx::query(&format!(
|
||||||
let mut where_clause = WhereClause::new();
|
"{SCORE_COLUMNS} WHERE id = ANY($1) ORDER BY id ASC"
|
||||||
push_in(&mut builder, &mut where_clause, "id", &query.ids);
|
))
|
||||||
let rows = builder
|
.bind(&query.ids)
|
||||||
.build()
|
.fetch_all(&self.pool)
|
||||||
.fetch_all(&self.pool)
|
.await
|
||||||
.await
|
.map_postgres_err()?;
|
||||||
.map_postgres_err()?;
|
|
||||||
rows.iter().map(map_score_row).collect()
|
rows.iter().map(map_score_row).collect()
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1,5 +1,5 @@
|
|||||||
use async_trait::async_trait;
|
use async_trait::async_trait;
|
||||||
use sqlx::{mysql::MySqlRow, Row};
|
use sqlx::{mysql::MySqlRow, MySql, QueryBuilder, Row};
|
||||||
|
|
||||||
use super::{
|
use super::{
|
||||||
InMemoryProviderCatalogReadRepository, ProviderCatalogKeyListQuery,
|
InMemoryProviderCatalogReadRepository, ProviderCatalogKeyListQuery,
|
||||||
@@ -16,6 +16,33 @@ pub struct MysqlProviderCatalogReadRepository {
|
|||||||
pool: MysqlPool,
|
pool: MysqlPool,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
const KEY_SELECT_SQL: &str = r#"
|
||||||
|
SELECT
|
||||||
|
id, provider_id, name, auth_type, capabilities, is_active, api_formats,
|
||||||
|
auth_type_by_format, allow_auth_channel_mismatch_formats,
|
||||||
|
COALESCE(api_key, encrypted_key) AS api_key,
|
||||||
|
auth_config, note, internal_priority, rate_multipliers,
|
||||||
|
global_priority_by_format, allowed_models,
|
||||||
|
expires_at AS expires_at_unix_secs,
|
||||||
|
cache_ttl_minutes, max_probe_interval_minutes, proxy, fingerprint,
|
||||||
|
rpm_limit, concurrent_limit, learned_rpm_limit, concurrent_429_count,
|
||||||
|
rpm_429_count, last_429_at AS last_429_at_unix_secs, last_429_type,
|
||||||
|
adjustment_history, utilization_samples,
|
||||||
|
last_probe_increase_at AS last_probe_increase_at_unix_secs,
|
||||||
|
last_rpm_peak, request_count, total_tokens, total_cost_usd,
|
||||||
|
success_count, error_count, total_response_time_ms,
|
||||||
|
last_used_at AS last_used_at_unix_secs, auto_fetch_models,
|
||||||
|
last_models_fetch_at AS last_models_fetch_at_unix_secs,
|
||||||
|
last_models_fetch_error, locked_models, model_include_patterns,
|
||||||
|
model_exclude_patterns, upstream_metadata,
|
||||||
|
oauth_invalid_at AS oauth_invalid_at_unix_secs,
|
||||||
|
oauth_invalid_reason, status_snapshot,
|
||||||
|
created_at AS created_at_unix_ms,
|
||||||
|
updated_at AS updated_at_unix_secs,
|
||||||
|
health_by_format, circuit_breaker_by_format
|
||||||
|
FROM provider_api_keys
|
||||||
|
"#;
|
||||||
|
|
||||||
impl MysqlProviderCatalogReadRepository {
|
impl MysqlProviderCatalogReadRepository {
|
||||||
pub fn new(pool: MysqlPool) -> Self {
|
pub fn new(pool: MysqlPool) -> Self {
|
||||||
Self { pool }
|
Self { pool }
|
||||||
@@ -71,37 +98,26 @@ WHERE api_format IS NOT NULL
|
|||||||
}
|
}
|
||||||
|
|
||||||
async fn load_keys(&self) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
|
async fn load_keys(&self) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
|
||||||
let rows = sqlx::query(
|
let rows = sqlx::query(KEY_SELECT_SQL)
|
||||||
r#"
|
.fetch_all(&self.pool)
|
||||||
SELECT
|
.await
|
||||||
id, provider_id, name, auth_type, capabilities, is_active, api_formats,
|
.map_sql_err()?;
|
||||||
auth_type_by_format, allow_auth_channel_mismatch_formats,
|
rows.iter().map(map_key_row).collect()
|
||||||
COALESCE(api_key, encrypted_key) AS api_key,
|
}
|
||||||
auth_config, note, internal_priority, rate_multipliers,
|
|
||||||
global_priority_by_format, allowed_models,
|
async fn list_keys_by_provider_ids_direct(
|
||||||
expires_at AS expires_at_unix_secs,
|
&self,
|
||||||
cache_ttl_minutes, max_probe_interval_minutes, proxy, fingerprint,
|
provider_ids: &[String],
|
||||||
rpm_limit, concurrent_limit, learned_rpm_limit, concurrent_429_count,
|
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
|
||||||
rpm_429_count, last_429_at AS last_429_at_unix_secs, last_429_type,
|
if provider_ids.is_empty() {
|
||||||
adjustment_history, utilization_samples,
|
return Ok(Vec::new());
|
||||||
last_probe_increase_at AS last_probe_increase_at_unix_secs,
|
}
|
||||||
last_rpm_peak, request_count, total_tokens, total_cost_usd,
|
|
||||||
success_count, error_count, total_response_time_ms,
|
let rows = build_list_keys_by_provider_ids_query(provider_ids)
|
||||||
last_used_at AS last_used_at_unix_secs, auto_fetch_models,
|
.build()
|
||||||
last_models_fetch_at AS last_models_fetch_at_unix_secs,
|
.fetch_all(&self.pool)
|
||||||
last_models_fetch_error, locked_models, model_include_patterns,
|
.await
|
||||||
model_exclude_patterns, upstream_metadata,
|
.map_sql_err()?;
|
||||||
oauth_invalid_at AS oauth_invalid_at_unix_secs,
|
|
||||||
oauth_invalid_reason, status_snapshot,
|
|
||||||
created_at AS created_at_unix_ms,
|
|
||||||
updated_at AS updated_at_unix_secs,
|
|
||||||
health_by_format, circuit_breaker_by_format
|
|
||||||
FROM provider_api_keys
|
|
||||||
"#,
|
|
||||||
)
|
|
||||||
.fetch_all(&self.pool)
|
|
||||||
.await
|
|
||||||
.map_sql_err()?;
|
|
||||||
rows.iter().map(map_key_row).collect()
|
rows.iter().map(map_key_row).collect()
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -912,20 +928,14 @@ impl ProviderCatalogReadRepository for MysqlProviderCatalogReadRepository {
|
|||||||
&self,
|
&self,
|
||||||
provider_ids: &[String],
|
provider_ids: &[String],
|
||||||
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
|
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
|
||||||
self.load_memory()
|
self.list_keys_by_provider_ids_direct(provider_ids).await
|
||||||
.await?
|
|
||||||
.list_keys_by_provider_ids(provider_ids)
|
|
||||||
.await
|
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn list_key_summaries_by_provider_ids(
|
async fn list_key_summaries_by_provider_ids(
|
||||||
&self,
|
&self,
|
||||||
provider_ids: &[String],
|
provider_ids: &[String],
|
||||||
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
|
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
|
||||||
self.load_memory()
|
self.list_keys_by_provider_ids_direct(provider_ids).await
|
||||||
.await?
|
|
||||||
.list_key_summaries_by_provider_ids(provider_ids)
|
|
||||||
.await
|
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn list_key_maintenance_summaries_by_provider_ids(
|
async fn list_key_maintenance_summaries_by_provider_ids(
|
||||||
@@ -1162,6 +1172,19 @@ fn optional_json_to_string(
|
|||||||
optional_json_ref_to_string(value.as_ref(), field_name)
|
optional_json_ref_to_string(value.as_ref(), field_name)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn build_list_keys_by_provider_ids_query(provider_ids: &[String]) -> QueryBuilder<'_, MySql> {
|
||||||
|
let mut builder = QueryBuilder::<MySql>::new(KEY_SELECT_SQL);
|
||||||
|
builder.push("WHERE provider_id IN (");
|
||||||
|
{
|
||||||
|
let mut separated = builder.separated(", ");
|
||||||
|
for provider_id in provider_ids {
|
||||||
|
separated.push_bind(provider_id.clone());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
builder.push(") ORDER BY provider_id ASC, name ASC, id ASC");
|
||||||
|
builder
|
||||||
|
}
|
||||||
|
|
||||||
fn key_insert_sql() -> &'static str {
|
fn key_insert_sql() -> &'static str {
|
||||||
r#"
|
r#"
|
||||||
INSERT INTO provider_api_keys (
|
INSERT INTO provider_api_keys (
|
||||||
@@ -1568,6 +1591,7 @@ mod tests {
|
|||||||
StoredProviderCatalogProvider,
|
StoredProviderCatalogProvider,
|
||||||
};
|
};
|
||||||
use serde_json::json;
|
use serde_json::json;
|
||||||
|
use sqlx::Execute;
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn repository_builds_from_lazy_pool() {
|
async fn repository_builds_from_lazy_pool() {
|
||||||
@@ -1580,6 +1604,17 @@ mod tests {
|
|||||||
let _repository = MysqlProviderCatalogReadRepository::new(pool);
|
let _repository = MysqlProviderCatalogReadRepository::new(pool);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn list_keys_by_provider_ids_query_targets_index_aligned_ordering() {
|
||||||
|
let provider_ids = vec!["provider-a".to_string(), "provider-b".to_string()];
|
||||||
|
let mut builder = super::build_list_keys_by_provider_ids_query(&provider_ids);
|
||||||
|
let query = builder.build();
|
||||||
|
let sql = query.sql();
|
||||||
|
|
||||||
|
assert!(sql.contains("WHERE provider_id IN ("));
|
||||||
|
assert!(sql.contains("ORDER BY provider_id ASC, name ASC, id ASC"));
|
||||||
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn mysql_provider_catalog_repository_round_trips_when_url_is_set() {
|
async fn mysql_provider_catalog_repository_round_trips_when_url_is_set() {
|
||||||
let Some(database_url) = std::env::var("AETHER_TEST_MYSQL_URL")
|
let Some(database_url) = std::env::var("AETHER_TEST_MYSQL_URL")
|
||||||
|
|||||||
+22
-18
@@ -114,15 +114,17 @@ SELECT
|
|||||||
OR NULLIF(BTRIM("usage".request_metadata->>'user_agent'), '') IS NOT NULL
|
OR NULLIF(BTRIM("usage".request_metadata->>'user_agent'), '') IS NOT NULL
|
||||||
OR NULLIF(BTRIM("usage".request_metadata->>'request_path'), '') IS NOT NULL
|
OR NULLIF(BTRIM("usage".request_metadata->>'request_path'), '') IS NOT NULL
|
||||||
OR NULLIF(BTRIM("usage".request_metadata->>'request_path_and_query'), '') IS NOT NULL
|
OR NULLIF(BTRIM("usage".request_metadata->>'request_path_and_query'), '') IS NOT NULL
|
||||||
OR COALESCE(
|
OR CASE
|
||||||
NULLIF(BTRIM("usage".request_metadata->>'provider_reasoning_effort'), ''),
|
WHEN jsonb_typeof("usage".provider_request_body::jsonb) = 'object' THEN COALESCE(
|
||||||
NULLIF(BTRIM("usage".provider_request_body->>'reasoning_effort'), ''),
|
NULLIF(BTRIM("usage".provider_request_body->>'reasoning_effort'), ''),
|
||||||
NULLIF(BTRIM("usage".provider_request_body->'reasoning'->>'effort'), '')
|
NULLIF(BTRIM("usage".provider_request_body->'reasoning'->>'effort'), '')
|
||||||
) IS NOT NULL
|
)
|
||||||
OR COALESCE(
|
ELSE NULLIF(BTRIM("usage".request_metadata->>'provider_reasoning_effort'), '')
|
||||||
NULLIF(BTRIM("usage".request_metadata->>'provider_service_tier'), ''),
|
END IS NOT NULL
|
||||||
NULLIF(BTRIM("usage".provider_request_body->>'service_tier'), '')
|
OR CASE
|
||||||
) IS NOT NULL
|
WHEN jsonb_typeof("usage".provider_request_body::jsonb) = 'object' THEN NULLIF(BTRIM("usage".provider_request_body->>'service_tier'), '')
|
||||||
|
ELSE NULLIF(BTRIM("usage".request_metadata->>'provider_service_tier'), '')
|
||||||
|
END IS NOT NULL
|
||||||
OR ("usage".request_metadata->>'client_requested_stream') IN ('true', 'false')
|
OR ("usage".request_metadata->>'client_requested_stream') IN ('true', 'false')
|
||||||
OR ("usage".request_metadata->>'upstream_is_stream') IN ('true', 'false')
|
OR ("usage".request_metadata->>'upstream_is_stream') IN ('true', 'false')
|
||||||
THEN jsonb_strip_nulls(jsonb_build_object(
|
THEN jsonb_strip_nulls(jsonb_build_object(
|
||||||
@@ -135,16 +137,18 @@ SELECT
|
|||||||
'request_path_and_query',
|
'request_path_and_query',
|
||||||
NULLIF(BTRIM("usage".request_metadata->>'request_path_and_query'), ''),
|
NULLIF(BTRIM("usage".request_metadata->>'request_path_and_query'), ''),
|
||||||
'provider_reasoning_effort',
|
'provider_reasoning_effort',
|
||||||
COALESCE(
|
CASE
|
||||||
NULLIF(BTRIM("usage".request_metadata->>'provider_reasoning_effort'), ''),
|
WHEN jsonb_typeof("usage".provider_request_body::jsonb) = 'object' THEN COALESCE(
|
||||||
NULLIF(BTRIM("usage".provider_request_body->>'reasoning_effort'), ''),
|
NULLIF(BTRIM("usage".provider_request_body->>'reasoning_effort'), ''),
|
||||||
NULLIF(BTRIM("usage".provider_request_body->'reasoning'->>'effort'), '')
|
NULLIF(BTRIM("usage".provider_request_body->'reasoning'->>'effort'), '')
|
||||||
),
|
)
|
||||||
|
ELSE NULLIF(BTRIM("usage".request_metadata->>'provider_reasoning_effort'), '')
|
||||||
|
END,
|
||||||
'provider_service_tier',
|
'provider_service_tier',
|
||||||
COALESCE(
|
CASE
|
||||||
NULLIF(BTRIM("usage".request_metadata->>'provider_service_tier'), ''),
|
WHEN jsonb_typeof("usage".provider_request_body::jsonb) = 'object' THEN NULLIF(BTRIM("usage".provider_request_body->>'service_tier'), '')
|
||||||
NULLIF(BTRIM("usage".provider_request_body->>'service_tier'), '')
|
ELSE NULLIF(BTRIM("usage".request_metadata->>'provider_service_tier'), '')
|
||||||
),
|
END,
|
||||||
'client_requested_stream',
|
'client_requested_stream',
|
||||||
CASE
|
CASE
|
||||||
WHEN ("usage".request_metadata->>'client_requested_stream') IN ('true', 'false')
|
WHEN ("usage".request_metadata->>'client_requested_stream') IN ('true', 'false')
|
||||||
|
|||||||
+22
-18
@@ -114,15 +114,17 @@ SELECT
|
|||||||
OR NULLIF(BTRIM("usage".request_metadata->>'user_agent'), '') IS NOT NULL
|
OR NULLIF(BTRIM("usage".request_metadata->>'user_agent'), '') IS NOT NULL
|
||||||
OR NULLIF(BTRIM("usage".request_metadata->>'request_path'), '') IS NOT NULL
|
OR NULLIF(BTRIM("usage".request_metadata->>'request_path'), '') IS NOT NULL
|
||||||
OR NULLIF(BTRIM("usage".request_metadata->>'request_path_and_query'), '') IS NOT NULL
|
OR NULLIF(BTRIM("usage".request_metadata->>'request_path_and_query'), '') IS NOT NULL
|
||||||
OR COALESCE(
|
OR CASE
|
||||||
NULLIF(BTRIM("usage".request_metadata->>'provider_reasoning_effort'), ''),
|
WHEN jsonb_typeof("usage".provider_request_body::jsonb) = 'object' THEN COALESCE(
|
||||||
NULLIF(BTRIM("usage".provider_request_body->>'reasoning_effort'), ''),
|
NULLIF(BTRIM("usage".provider_request_body->>'reasoning_effort'), ''),
|
||||||
NULLIF(BTRIM("usage".provider_request_body->'reasoning'->>'effort'), '')
|
NULLIF(BTRIM("usage".provider_request_body->'reasoning'->>'effort'), '')
|
||||||
) IS NOT NULL
|
)
|
||||||
OR COALESCE(
|
ELSE NULLIF(BTRIM("usage".request_metadata->>'provider_reasoning_effort'), '')
|
||||||
NULLIF(BTRIM("usage".request_metadata->>'provider_service_tier'), ''),
|
END IS NOT NULL
|
||||||
NULLIF(BTRIM("usage".provider_request_body->>'service_tier'), '')
|
OR CASE
|
||||||
) IS NOT NULL
|
WHEN jsonb_typeof("usage".provider_request_body::jsonb) = 'object' THEN NULLIF(BTRIM("usage".provider_request_body->>'service_tier'), '')
|
||||||
|
ELSE NULLIF(BTRIM("usage".request_metadata->>'provider_service_tier'), '')
|
||||||
|
END IS NOT NULL
|
||||||
OR ("usage".request_metadata->>'client_requested_stream') IN ('true', 'false')
|
OR ("usage".request_metadata->>'client_requested_stream') IN ('true', 'false')
|
||||||
OR ("usage".request_metadata->>'upstream_is_stream') IN ('true', 'false')
|
OR ("usage".request_metadata->>'upstream_is_stream') IN ('true', 'false')
|
||||||
THEN jsonb_strip_nulls(jsonb_build_object(
|
THEN jsonb_strip_nulls(jsonb_build_object(
|
||||||
@@ -135,16 +137,18 @@ SELECT
|
|||||||
'request_path_and_query',
|
'request_path_and_query',
|
||||||
NULLIF(BTRIM("usage".request_metadata->>'request_path_and_query'), ''),
|
NULLIF(BTRIM("usage".request_metadata->>'request_path_and_query'), ''),
|
||||||
'provider_reasoning_effort',
|
'provider_reasoning_effort',
|
||||||
COALESCE(
|
CASE
|
||||||
NULLIF(BTRIM("usage".request_metadata->>'provider_reasoning_effort'), ''),
|
WHEN jsonb_typeof("usage".provider_request_body::jsonb) = 'object' THEN COALESCE(
|
||||||
NULLIF(BTRIM("usage".provider_request_body->>'reasoning_effort'), ''),
|
NULLIF(BTRIM("usage".provider_request_body->>'reasoning_effort'), ''),
|
||||||
NULLIF(BTRIM("usage".provider_request_body->'reasoning'->>'effort'), '')
|
NULLIF(BTRIM("usage".provider_request_body->'reasoning'->>'effort'), '')
|
||||||
),
|
)
|
||||||
|
ELSE NULLIF(BTRIM("usage".request_metadata->>'provider_reasoning_effort'), '')
|
||||||
|
END,
|
||||||
'provider_service_tier',
|
'provider_service_tier',
|
||||||
COALESCE(
|
CASE
|
||||||
NULLIF(BTRIM("usage".request_metadata->>'provider_service_tier'), ''),
|
WHEN jsonb_typeof("usage".provider_request_body::jsonb) = 'object' THEN NULLIF(BTRIM("usage".provider_request_body->>'service_tier'), '')
|
||||||
NULLIF(BTRIM("usage".provider_request_body->>'service_tier'), '')
|
ELSE NULLIF(BTRIM("usage".request_metadata->>'provider_service_tier'), '')
|
||||||
),
|
END,
|
||||||
'client_requested_stream',
|
'client_requested_stream',
|
||||||
CASE
|
CASE
|
||||||
WHEN ("usage".request_metadata->>'client_requested_stream') IN ('true', 'false')
|
WHEN ("usage".request_metadata->>'client_requested_stream') IN ('true', 'false')
|
||||||
|
|||||||
@@ -707,6 +707,8 @@ fn usage_sql_uses_json_null_placeholders_for_usage_payload_columns() {
|
|||||||
super::LIST_RECENT_USAGE_AUDITS_PREFIX,
|
super::LIST_RECENT_USAGE_AUDITS_PREFIX,
|
||||||
] {
|
] {
|
||||||
assert!(sql.contains("jsonb_strip_nulls(jsonb_build_object("));
|
assert!(sql.contains("jsonb_strip_nulls(jsonb_build_object("));
|
||||||
|
assert!(sql.contains("jsonb_typeof(\"usage\".provider_request_body::jsonb)"));
|
||||||
|
assert!(!sql.contains("jsonb_typeof(\"usage\".provider_request_body)"));
|
||||||
assert!(sql.contains("'client_ip'"));
|
assert!(sql.contains("'client_ip'"));
|
||||||
assert!(sql.contains("request_metadata->>'client_ip'"));
|
assert!(sql.contains("request_metadata->>'client_ip'"));
|
||||||
assert!(sql.contains("'user_agent'"));
|
assert!(sql.contains("'user_agent'"));
|
||||||
|
|||||||
@@ -548,6 +548,24 @@ mod tests {
|
|||||||
.as_deref(),
|
.as_deref(),
|
||||||
Some("gpt-5.4")
|
Some("gpt-5.4")
|
||||||
);
|
);
|
||||||
|
|
||||||
|
let row = sample_row("gpt-5.4", "gpt-5.4-upstream");
|
||||||
|
assert!(row_supports_requested_model_with_model_directives(
|
||||||
|
&row,
|
||||||
|
"gpt-5.4-fast-xhigh",
|
||||||
|
"openai:chat",
|
||||||
|
true
|
||||||
|
));
|
||||||
|
assert_eq!(
|
||||||
|
resolve_requested_global_model_name_with_model_directives(
|
||||||
|
&[row],
|
||||||
|
"gpt-5.4-xhigh-fast",
|
||||||
|
"openai:chat",
|
||||||
|
true
|
||||||
|
)
|
||||||
|
.as_deref(),
|
||||||
|
Some("gpt-5.4")
|
||||||
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
|
|||||||
@@ -77,15 +77,20 @@ pub(crate) fn attach_provider_request_body_metadata(
|
|||||||
metadata: Option<Value>,
|
metadata: Option<Value>,
|
||||||
provider_request_body: Option<&Value>,
|
provider_request_body: Option<&Value>,
|
||||||
) -> Option<Value> {
|
) -> Option<Value> {
|
||||||
|
let provider_body_is_object = provider_request_body.and_then(Value::as_object).is_some();
|
||||||
let reasoning_effort = extract_provider_reasoning_effort_from_body(provider_request_body);
|
let reasoning_effort = extract_provider_reasoning_effort_from_body(provider_request_body);
|
||||||
let service_tier = extract_provider_service_tier_from_body(provider_request_body);
|
let service_tier = extract_provider_service_tier_from_body(provider_request_body);
|
||||||
if reasoning_effort.is_none() && service_tier.is_none() {
|
if !provider_body_is_object && reasoning_effort.is_none() && service_tier.is_none() {
|
||||||
return metadata;
|
return metadata;
|
||||||
}
|
}
|
||||||
let mut object = match metadata {
|
let mut object = match metadata {
|
||||||
Some(Value::Object(object)) => object,
|
Some(Value::Object(object)) => object,
|
||||||
_ => Map::new(),
|
_ => Map::new(),
|
||||||
};
|
};
|
||||||
|
if provider_body_is_object {
|
||||||
|
object.remove(PROVIDER_REASONING_EFFORT_METADATA_KEY);
|
||||||
|
object.remove(PROVIDER_SERVICE_TIER_METADATA_KEY);
|
||||||
|
}
|
||||||
if let Some(reasoning_effort) = reasoning_effort {
|
if let Some(reasoning_effort) = reasoning_effort {
|
||||||
object.insert(
|
object.insert(
|
||||||
PROVIDER_REASONING_EFFORT_METADATA_KEY.to_string(),
|
PROVIDER_REASONING_EFFORT_METADATA_KEY.to_string(),
|
||||||
@@ -98,7 +103,7 @@ pub(crate) fn attach_provider_request_body_metadata(
|
|||||||
Value::String(service_tier),
|
Value::String(service_tier),
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
Some(Value::Object(object))
|
(!object.is_empty()).then_some(Value::Object(object))
|
||||||
}
|
}
|
||||||
|
|
||||||
fn copy_allowed_metadata_fields(source: &Map<String, Value>, target: &mut Map<String, Value>) {
|
fn copy_allowed_metadata_fields(source: &Map<String, Value>, target: &mut Map<String, Value>) {
|
||||||
@@ -438,10 +443,11 @@ mod tests {
|
|||||||
use std::collections::BTreeMap;
|
use std::collections::BTreeMap;
|
||||||
|
|
||||||
use super::{
|
use super::{
|
||||||
build_usage_request_metadata_seed, merge_usage_request_metadata,
|
attach_provider_request_body_metadata, build_usage_request_metadata_seed,
|
||||||
merge_usage_request_metadata_owned, sanitize_usage_request_metadata,
|
merge_usage_request_metadata, merge_usage_request_metadata_owned,
|
||||||
sanitize_usage_request_metadata_ref, MAX_USAGE_REQUEST_METADATA_BYTES,
|
sanitize_usage_request_metadata, sanitize_usage_request_metadata_ref,
|
||||||
MAX_USAGE_REQUEST_METADATA_DEPTH, MAX_USAGE_REQUEST_METADATA_NODES,
|
MAX_USAGE_REQUEST_METADATA_BYTES, MAX_USAGE_REQUEST_METADATA_DEPTH,
|
||||||
|
MAX_USAGE_REQUEST_METADATA_NODES,
|
||||||
};
|
};
|
||||||
|
|
||||||
fn sample_plan() -> ExecutionPlan {
|
fn sample_plan() -> ExecutionPlan {
|
||||||
@@ -678,6 +684,49 @@ mod tests {
|
|||||||
assert_eq!(metadata, None);
|
assert_eq!(metadata, None);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn provider_request_body_metadata_uses_final_provider_body_as_source_of_truth() {
|
||||||
|
let metadata = Some(json!({
|
||||||
|
"trace_id": "trace-1",
|
||||||
|
"provider_reasoning_effort": "high",
|
||||||
|
"provider_service_tier": "priority"
|
||||||
|
}));
|
||||||
|
|
||||||
|
let updated = attach_provider_request_body_metadata(
|
||||||
|
metadata.clone(),
|
||||||
|
Some(&json!({
|
||||||
|
"model": "gpt-5",
|
||||||
|
"reasoning": { "effort": "low" },
|
||||||
|
"service_tier": "standard"
|
||||||
|
})),
|
||||||
|
)
|
||||||
|
.expect("metadata should remain");
|
||||||
|
|
||||||
|
assert_eq!(
|
||||||
|
updated,
|
||||||
|
json!({
|
||||||
|
"trace_id": "trace-1",
|
||||||
|
"provider_reasoning_effort": "low",
|
||||||
|
"provider_service_tier": "standard"
|
||||||
|
})
|
||||||
|
);
|
||||||
|
|
||||||
|
let cleared = attach_provider_request_body_metadata(
|
||||||
|
metadata,
|
||||||
|
Some(&json!({
|
||||||
|
"model": "gpt-5"
|
||||||
|
})),
|
||||||
|
)
|
||||||
|
.expect("metadata should retain unrelated fields");
|
||||||
|
|
||||||
|
assert_eq!(
|
||||||
|
cleared,
|
||||||
|
json!({
|
||||||
|
"trace_id": "trace-1"
|
||||||
|
})
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn owned_merge_matches_filtered_merge_for_trusted_objects() {
|
fn owned_merge_matches_filtered_merge_for_trusted_objects() {
|
||||||
let base = Some(json!({
|
let base = Some(json!({
|
||||||
|
|||||||
@@ -1,12 +1,14 @@
|
|||||||
import { beforeEach, describe, expect, it, vi } from 'vitest'
|
import { beforeEach, describe, expect, it, vi } from 'vitest'
|
||||||
|
|
||||||
const { postMock } = vi.hoisted(() => ({
|
const { postMock, setTokenMock } = vi.hoisted(() => ({
|
||||||
postMock: vi.fn(),
|
postMock: vi.fn(),
|
||||||
|
setTokenMock: vi.fn(),
|
||||||
}))
|
}))
|
||||||
|
|
||||||
vi.mock('@/api/client', () => ({
|
vi.mock('@/api/client', () => ({
|
||||||
default: {
|
default: {
|
||||||
post: postMock,
|
post: postMock,
|
||||||
|
setToken: setTokenMock,
|
||||||
},
|
},
|
||||||
}))
|
}))
|
||||||
|
|
||||||
@@ -15,6 +17,7 @@ import { authApi } from '@/api/auth'
|
|||||||
describe('authApi turnstile payloads', () => {
|
describe('authApi turnstile payloads', () => {
|
||||||
beforeEach(() => {
|
beforeEach(() => {
|
||||||
postMock.mockReset()
|
postMock.mockReset()
|
||||||
|
setTokenMock.mockReset()
|
||||||
postMock.mockResolvedValue({ data: {} })
|
postMock.mockResolvedValue({ data: {} })
|
||||||
})
|
})
|
||||||
|
|
||||||
@@ -42,4 +45,13 @@ describe('authApi turnstile payloads', () => {
|
|||||||
turnstile_token: 'turnstile-token',
|
turnstile_token: 'turnstile-token',
|
||||||
})
|
})
|
||||||
})
|
})
|
||||||
|
|
||||||
|
it('refreshes auth token without a request body', async () => {
|
||||||
|
postMock.mockResolvedValue({ data: { access_token: 'new-access-token' } })
|
||||||
|
|
||||||
|
await authApi.refreshToken()
|
||||||
|
|
||||||
|
expect(postMock).toHaveBeenCalledWith('/api/auth/refresh')
|
||||||
|
expect(setTokenMock).toHaveBeenCalledWith('new-access-token')
|
||||||
|
})
|
||||||
})
|
})
|
||||||
|
|||||||
@@ -1,7 +1,12 @@
|
|||||||
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
|
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
|
||||||
|
import type { AxiosAdapter, AxiosInstance, InternalAxiosRequestConfig } from 'axios'
|
||||||
|
|
||||||
import apiClient, { AUTH_STATE_CHANGE_EVENT } from '@/api/client'
|
import apiClient, { AUTH_STATE_CHANGE_EVENT } from '@/api/client'
|
||||||
|
|
||||||
|
type TestableApiClient = typeof apiClient & {
|
||||||
|
client: AxiosInstance
|
||||||
|
}
|
||||||
|
|
||||||
describe('apiClient auth state change event', () => {
|
describe('apiClient auth state change event', () => {
|
||||||
beforeEach(() => {
|
beforeEach(() => {
|
||||||
localStorage.clear()
|
localStorage.clear()
|
||||||
@@ -28,4 +33,33 @@ describe('apiClient auth state change event', () => {
|
|||||||
|
|
||||||
window.removeEventListener(AUTH_STATE_CHANGE_EVENT, handler as EventListener)
|
window.removeEventListener(AUTH_STATE_CHANGE_EVENT, handler as EventListener)
|
||||||
})
|
})
|
||||||
|
|
||||||
|
it('sends auth refresh without a request body', async () => {
|
||||||
|
const rawClient = apiClient as TestableApiClient
|
||||||
|
const previousAdapter = rawClient.client.defaults.adapter
|
||||||
|
const requests: InternalAxiosRequestConfig[] = []
|
||||||
|
|
||||||
|
rawClient.client.defaults.adapter = (async (config: InternalAxiosRequestConfig) => {
|
||||||
|
requests.push(config)
|
||||||
|
return {
|
||||||
|
data: { access_token: 'new-access-token' },
|
||||||
|
status: 200,
|
||||||
|
statusText: 'OK',
|
||||||
|
headers: {},
|
||||||
|
config,
|
||||||
|
}
|
||||||
|
}) as AxiosAdapter
|
||||||
|
|
||||||
|
try {
|
||||||
|
const response = await apiClient.refreshToken()
|
||||||
|
|
||||||
|
expect(response.data.access_token).toBe('new-access-token')
|
||||||
|
expect(requests).toHaveLength(1)
|
||||||
|
expect(requests[0].url).toBe('/api/auth/refresh')
|
||||||
|
expect(requests[0].method).toBe('post')
|
||||||
|
expect(requests[0].data).toBeUndefined()
|
||||||
|
} finally {
|
||||||
|
rawClient.client.defaults.adapter = previousAdapter
|
||||||
|
}
|
||||||
|
})
|
||||||
})
|
})
|
||||||
|
|||||||
@@ -164,7 +164,7 @@ export const authApi = {
|
|||||||
},
|
},
|
||||||
|
|
||||||
async refreshToken(): Promise<LoginResponse> {
|
async refreshToken(): Promise<LoginResponse> {
|
||||||
const response = await apiClient.post<LoginResponse>('/api/auth/refresh', {})
|
const response = await apiClient.post<LoginResponse>('/api/auth/refresh')
|
||||||
apiClient.setToken(response.data.access_token)
|
apiClient.setToken(response.data.access_token)
|
||||||
return response.data
|
return response.data
|
||||||
},
|
},
|
||||||
|
|||||||
@@ -323,7 +323,7 @@ class ApiClient {
|
|||||||
}
|
}
|
||||||
|
|
||||||
async refreshToken(): Promise<AxiosResponse> {
|
async refreshToken(): Promise<AxiosResponse> {
|
||||||
return this.client.post('/api/auth/refresh', {})
|
return this.client.post('/api/auth/refresh')
|
||||||
}
|
}
|
||||||
|
|
||||||
// 以下方法直接委托给 axios client,Demo 模式由 adapter 统一处理
|
// 以下方法直接委托给 axios client,Demo 模式由 adapter 统一处理
|
||||||
|
|||||||
@@ -203,83 +203,33 @@
|
|||||||
<!-- 费用与性能概览 -->
|
<!-- 费用与性能概览 -->
|
||||||
<Card>
|
<Card>
|
||||||
<div class="p-3 sm:p-4">
|
<div class="p-3 sm:p-4">
|
||||||
<!-- 总费用和响应时间(独立显示) -->
|
<div class="mb-4 flex flex-wrap items-center gap-x-2 gap-y-1 text-sm">
|
||||||
<div
|
<span class="rounded bg-muted px-1.5 py-0.5 text-[10px] text-muted-foreground/70">{{ priceSourceLabel }}</span>
|
||||||
class="grid gap-1.5 mb-4 sm:gap-2"
|
<span class="text-muted-foreground">|</span>
|
||||||
:class="detailOutputRate != null ? 'grid-cols-3' : 'grid-cols-2'"
|
<span class="whitespace-nowrap">
|
||||||
>
|
<span class="text-muted-foreground">总费用</span>
|
||||||
<div class="min-w-0 rounded-md border border-border/50 bg-muted/20 px-2 py-2 sm:px-3">
|
<span class="ml-1 font-bold text-green-600 dark:text-green-400">
|
||||||
<span class="mb-0.5 block truncate text-[10px] text-muted-foreground sm:text-xs">总费用</span>
|
|
||||||
<span class="block truncate text-sm font-bold text-green-600 dark:text-green-400 sm:text-lg">
|
|
||||||
${{ ((typeof detail.cost === 'object' ? detail.cost?.total : detail.cost) || detail.total_cost || 0).toFixed(6) }}
|
${{ ((typeof detail.cost === 'object' ? detail.cost?.total : detail.cost) || detail.total_cost || 0).toFixed(6) }}
|
||||||
</span>
|
</span>
|
||||||
</div>
|
</span>
|
||||||
<div class="min-w-0 rounded-md border border-border/50 bg-muted/20 px-2 py-2 sm:px-3">
|
<span class="text-muted-foreground">|</span>
|
||||||
<span class="mb-0.5 block truncate text-[10px] text-muted-foreground sm:text-xs">响应时间</span>
|
<span class="whitespace-nowrap">
|
||||||
<span class="text-sm font-bold sm:text-lg">{{ detail.response_time_ms ? formatResponseTime(detail.response_time_ms).value : 'N/A' }}</span>
|
<span class="text-muted-foreground">耗时</span>
|
||||||
<span class="ml-0.5 text-[10px] text-muted-foreground sm:ml-1 sm:text-sm">{{ detail.response_time_ms ? formatResponseTime(detail.response_time_ms).unit : '' }}</span>
|
<span class="ml-1 font-bold">
|
||||||
</div>
|
{{ formatDurationMs(detail.first_byte_time_ms) }} / {{ formatDurationMs(detail.response_time_ms) }}
|
||||||
<template v-if="detailOutputRate != null">
|
</span>
|
||||||
<div class="min-w-0 rounded-md border border-border/50 bg-muted/20 px-2 py-2 sm:px-3">
|
</span>
|
||||||
<span class="mb-0.5 block truncate text-[10px] text-muted-foreground sm:text-xs">输出速度</span>
|
<span class="text-muted-foreground">|</span>
|
||||||
<span class="text-sm font-bold text-primary sm:text-lg">{{ formatOutputRateValue(detailOutputRate) }}</span>
|
<span class="whitespace-nowrap">
|
||||||
<span class="ml-0.5 text-[10px] text-muted-foreground sm:ml-1 sm:text-sm">tps</span>
|
<span class="text-muted-foreground">输出速度</span>
|
||||||
</div>
|
<span class="ml-1 font-bold text-primary">{{ formatOutputRateValue(detailOutputRate) }}tps</span>
|
||||||
</template>
|
</span>
|
||||||
</div>
|
|
||||||
|
|
||||||
<div
|
|
||||||
v-if="hasDetailPerformanceBreakdown"
|
|
||||||
class="grid grid-cols-3 gap-1.5 mb-4 text-xs sm:gap-2"
|
|
||||||
>
|
|
||||||
<div class="min-w-0 rounded-md border border-border/50 bg-muted/20 px-2 py-2 sm:px-3">
|
|
||||||
<div class="mb-1 truncate text-[10px] text-muted-foreground sm:text-xs">
|
|
||||||
首字时间
|
|
||||||
</div>
|
|
||||||
<div class="truncate font-mono text-xs text-foreground sm:text-sm">
|
|
||||||
{{ formatDurationMs(detail.first_byte_time_ms) }}
|
|
||||||
</div>
|
|
||||||
</div>
|
|
||||||
<div class="min-w-0 rounded-md border border-border/50 bg-muted/20 px-2 py-2 sm:px-3">
|
|
||||||
<div class="mb-1 truncate text-[10px] text-muted-foreground sm:text-xs">
|
|
||||||
生成耗时
|
|
||||||
</div>
|
|
||||||
<div class="truncate font-mono text-xs text-foreground sm:text-sm">
|
|
||||||
{{ formatDurationMs(detailGenerationTimeMs) }}
|
|
||||||
</div>
|
|
||||||
</div>
|
|
||||||
<div class="min-w-0 rounded-md border border-border/50 bg-muted/20 px-2 py-2 sm:px-3">
|
|
||||||
<div class="mb-1 truncate text-[10px] text-muted-foreground sm:text-xs">
|
|
||||||
输出 Tokens
|
|
||||||
</div>
|
|
||||||
<div class="truncate font-mono text-xs text-foreground sm:text-sm">
|
|
||||||
{{ formatNumber(detailOutputTokens) }}
|
|
||||||
</div>
|
|
||||||
</div>
|
|
||||||
</div>
|
</div>
|
||||||
|
|
||||||
<!-- 分隔线 -->
|
<!-- 分隔线 -->
|
||||||
<Separator class="mb-4" />
|
<Separator class="mb-4" />
|
||||||
|
|
||||||
<!-- ========== 1. 费用聚合计算 ========== -->
|
<!-- ========== 1. Token分阶段成本 ========== -->
|
||||||
<div class="text-xs text-muted-foreground mb-3 flex items-center gap-2 flex-wrap">
|
|
||||||
<span class="text-[10px] px-1.5 py-0.5 rounded bg-muted text-muted-foreground/70">{{ priceSourceLabel }}</span>
|
|
||||||
<span class="text-foreground">|</span>
|
|
||||||
<span class="font-mono text-foreground">
|
|
||||||
总费用 = Token费用 <span class="font-medium">${{ tokenCostTotal.toFixed(6) }}</span>
|
|
||||||
<template v-if="perRequestCost > 0">
|
|
||||||
+ 按次费用 <span class="font-medium">${{ perRequestCost.toFixed(6) }}</span>
|
|
||||||
</template>
|
|
||||||
<template v-if="imageOutputCostTotal > 0">
|
|
||||||
+ 图片输出费用 <span class="font-medium">${{ imageOutputCostTotal.toFixed(6) }}</span>
|
|
||||||
</template>
|
|
||||||
<template v-if="videoCostTotal > 0">
|
|
||||||
+ {{ detail.video_billing?.task_type === 'image' ? '图像' : detail.video_billing?.task_type === 'audio' ? '音频' : '视频' }}费用 <span class="font-medium">${{ videoCostTotal.toFixed(6) }}</span>
|
|
||||||
</template>
|
|
||||||
</span>
|
|
||||||
</div>
|
|
||||||
|
|
||||||
<!-- ========== 2. Token分阶段成本 ========== -->
|
|
||||||
<div
|
<div
|
||||||
v-if="hasTokenCost"
|
v-if="hasTokenCost"
|
||||||
class="space-y-2 mb-3"
|
class="space-y-2 mb-3"
|
||||||
@@ -287,6 +237,7 @@
|
|||||||
<!-- 阶梯标题 -->
|
<!-- 阶梯标题 -->
|
||||||
<div class="text-xs text-muted-foreground flex items-center gap-2 flex-wrap">
|
<div class="text-xs text-muted-foreground flex items-center gap-2 flex-wrap">
|
||||||
<span class="font-medium text-foreground">Token 计费</span>
|
<span class="font-medium text-foreground">Token 计费</span>
|
||||||
|
<span class="font-mono font-medium text-foreground">${{ tokenCostTotal.toFixed(6) }}</span>
|
||||||
<span class="text-muted-foreground/60">(输入 {{ formatNumber(displayInputTokens) }} + 缓存创建 {{ cacheCreationSummaryText }} + 缓存读取 {{ formatNumber(detail.cache_read_input_tokens || 0) }})</span>
|
<span class="text-muted-foreground/60">(输入 {{ formatNumber(displayInputTokens) }} + 缓存创建 {{ cacheCreationSummaryText }} + 缓存读取 {{ formatNumber(detail.cache_read_input_tokens || 0) }})</span>
|
||||||
<Badge
|
<Badge
|
||||||
v-if="displayTiers.length > 1"
|
v-if="displayTiers.length > 1"
|
||||||
@@ -404,8 +355,9 @@
|
|||||||
v-if="perRequestCost > 0 && !detail.video_billing"
|
v-if="perRequestCost > 0 && !detail.video_billing"
|
||||||
class="space-y-2 mb-3"
|
class="space-y-2 mb-3"
|
||||||
>
|
>
|
||||||
<div class="flex items-center justify-between text-xs">
|
<div class="flex items-center gap-2 text-xs">
|
||||||
<span class="font-medium text-foreground">按次计费</span>
|
<span class="font-medium text-foreground">按次计费</span>
|
||||||
|
<span class="font-mono font-medium text-foreground">${{ perRequestCost.toFixed(6) }}</span>
|
||||||
</div>
|
</div>
|
||||||
<div class="rounded-lg p-3 bg-primary/5 border border-primary/30 space-y-2">
|
<div class="rounded-lg p-3 bg-primary/5 border border-primary/30 space-y-2">
|
||||||
<div
|
<div
|
||||||
@@ -797,7 +749,6 @@ import {
|
|||||||
formatOutputRate,
|
formatOutputRate,
|
||||||
formatOutputRateValue,
|
formatOutputRateValue,
|
||||||
getDisplayOutputRate,
|
getDisplayOutputRate,
|
||||||
getGenerationTimeMs,
|
|
||||||
} from '../performance'
|
} from '../performance'
|
||||||
import {
|
import {
|
||||||
formatUsageStreamLabel,
|
formatUsageStreamLabel,
|
||||||
@@ -1108,16 +1059,6 @@ const detailOutputTokens = computed(() => {
|
|||||||
return detail.value.tokens?.output ?? detail.value.output_tokens ?? 0
|
return detail.value.tokens?.output ?? detail.value.output_tokens ?? 0
|
||||||
})
|
})
|
||||||
|
|
||||||
const detailGenerationTimeMs = computed(() => {
|
|
||||||
if (!detail.value) return null
|
|
||||||
return getGenerationTimeMs({
|
|
||||||
response_time_ms: detail.value.response_time_ms,
|
|
||||||
first_byte_time_ms: detail.value.first_byte_time_ms,
|
|
||||||
is_stream: detail.value.is_stream,
|
|
||||||
upstream_is_stream: detail.value.upstream_is_stream,
|
|
||||||
})
|
|
||||||
})
|
|
||||||
|
|
||||||
const detailOutputRate = computed(() => {
|
const detailOutputRate = computed(() => {
|
||||||
if (!detail.value) return null
|
if (!detail.value) return null
|
||||||
return getDisplayOutputRate({
|
return getDisplayOutputRate({
|
||||||
@@ -1129,11 +1070,6 @@ const detailOutputRate = computed(() => {
|
|||||||
})
|
})
|
||||||
})
|
})
|
||||||
|
|
||||||
const hasDetailPerformanceBreakdown = computed(() => {
|
|
||||||
if (!detail.value) return false
|
|
||||||
return detail.value.first_byte_time_ms != null || detailGenerationTimeMs.value != null || detailOutputTokens.value > 0
|
|
||||||
})
|
|
||||||
|
|
||||||
// 监听标签页切换
|
// 监听标签页切换
|
||||||
watch(activeTab, (newTab) => {
|
watch(activeTab, (newTab) => {
|
||||||
if (!['request-headers', 'response-headers'].includes(newTab) && viewMode.value === 'compare') {
|
if (!['request-headers', 'response-headers'].includes(newTab) && viewMode.value === 'compare') {
|
||||||
@@ -2378,14 +2314,6 @@ function formatPixels(value: number): string {
|
|||||||
return `${formatCompactNumber(value)} px`
|
return `${formatCompactNumber(value)} px`
|
||||||
}
|
}
|
||||||
|
|
||||||
// 格式化响应时间,自动选择合适的单位
|
|
||||||
function formatResponseTime(ms: number): { value: string; unit: string } {
|
|
||||||
if (ms >= 1_000) {
|
|
||||||
return { value: (ms / 1_000).toFixed(2), unit: 's' }
|
|
||||||
}
|
|
||||||
return { value: ms.toString(), unit: 'ms' }
|
|
||||||
}
|
|
||||||
|
|
||||||
// 格式化价格,修复浮点数精度问题
|
// 格式化价格,修复浮点数精度问题
|
||||||
function formatPrice(price: number): string {
|
function formatPrice(price: number): string {
|
||||||
// 处理浮点数精度问题,最多保留4位小数,去掉尾部的0
|
// 处理浮点数精度问题,最多保留4位小数,去掉尾部的0
|
||||||
|
|||||||
Reference in New Issue
Block a user