mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-03 09:50:21 +08:00
Merge remote-tracking branch 'origin/codex/pool-member-scores' into aether-rust-pioneer
This commit is contained in:
1
Cargo.lock
generated
1
Cargo.lock
generated
@@ -63,6 +63,7 @@ version = "0.1.0"
|
|||||||
dependencies = [
|
dependencies = [
|
||||||
"aether-ai-formats",
|
"aether-ai-formats",
|
||||||
"aether-contracts",
|
"aether-contracts",
|
||||||
|
"aether-data-contracts",
|
||||||
"aether-scheduler-core",
|
"aether-scheduler-core",
|
||||||
"async-trait",
|
"async-trait",
|
||||||
"http",
|
"http",
|
||||||
|
|||||||
@@ -43,13 +43,13 @@ pub(crate) use self::planner::{
|
|||||||
build_local_video_sync_plan_and_reports_for_kind,
|
build_local_video_sync_plan_and_reports_for_kind,
|
||||||
build_openai_responses_stream_plan_from_decision,
|
build_openai_responses_stream_plan_from_decision,
|
||||||
build_openai_responses_sync_plan_from_decision, build_passthrough_sync_plan_from_decision,
|
build_openai_responses_sync_plan_from_decision, build_passthrough_sync_plan_from_decision,
|
||||||
build_standard_family_stream_attempt_source, build_standard_family_stream_plan_and_reports,
|
build_provider_key_pool_score_upsert, build_standard_family_stream_attempt_source,
|
||||||
build_standard_family_sync_attempt_source, build_standard_family_sync_plan_and_reports,
|
build_standard_family_stream_plan_and_reports, build_standard_family_sync_attempt_source,
|
||||||
build_standard_stream_plan_from_decision, build_standard_sync_plan_from_decision,
|
build_standard_family_sync_plan_and_reports, build_standard_stream_plan_from_decision,
|
||||||
extract_pool_sticky_session_token, maybe_build_stream_decision_payload,
|
build_standard_sync_plan_from_decision, extract_pool_sticky_session_token,
|
||||||
maybe_build_stream_plan_payload, maybe_build_sync_decision_payload,
|
maybe_build_stream_decision_payload, maybe_build_stream_plan_payload,
|
||||||
maybe_build_sync_plan_payload, planner_is_matching_stream_request,
|
maybe_build_sync_decision_payload, maybe_build_sync_plan_payload,
|
||||||
set_local_openai_chat_execution_exhausted_diagnostic,
|
planner_is_matching_stream_request, set_local_openai_chat_execution_exhausted_diagnostic,
|
||||||
set_local_openai_image_execution_exhausted_diagnostic, CandidateFailureDiagnostic,
|
set_local_openai_image_execution_exhausted_diagnostic, CandidateFailureDiagnostic,
|
||||||
CandidateFailureDiagnosticKind, GatewayAuthApiKeySnapshot, GatewayProviderTransportSnapshot,
|
CandidateFailureDiagnosticKind, GatewayAuthApiKeySnapshot, GatewayProviderTransportSnapshot,
|
||||||
LocalExecutionAttemptSource, LocalResolvedOAuthRequestAuth, PlannerAppState,
|
LocalExecutionAttemptSource, LocalResolvedOAuthRequestAuth, PlannerAppState,
|
||||||
|
|||||||
@@ -16,6 +16,7 @@ mod materialization_policy;
|
|||||||
mod passthrough;
|
mod passthrough;
|
||||||
mod plan_builders;
|
mod plan_builders;
|
||||||
mod pool_scheduler;
|
mod pool_scheduler;
|
||||||
|
pub(crate) mod pool_scores;
|
||||||
mod report_context;
|
mod report_context;
|
||||||
mod route;
|
mod route;
|
||||||
mod runtime_miss;
|
mod runtime_miss;
|
||||||
@@ -36,6 +37,7 @@ pub(crate) use self::plan_builders::{
|
|||||||
build_standard_stream_plan_from_decision, build_standard_sync_plan_from_decision,
|
build_standard_stream_plan_from_decision, build_standard_sync_plan_from_decision,
|
||||||
AiStreamAttempt, AiSyncAttempt,
|
AiStreamAttempt, AiSyncAttempt,
|
||||||
};
|
};
|
||||||
|
pub(crate) use self::pool_scores::build_provider_key_pool_score_upsert;
|
||||||
pub(crate) use self::route::is_matching_stream_request as planner_is_matching_stream_request;
|
pub(crate) use self::route::is_matching_stream_request as planner_is_matching_stream_request;
|
||||||
pub(crate) use self::specialized::{
|
pub(crate) use self::specialized::{
|
||||||
build_local_gemini_files_stream_attempt_source_for_kind,
|
build_local_gemini_files_stream_attempt_source_for_kind,
|
||||||
|
|||||||
@@ -9,7 +9,11 @@ use aether_ai_serving::{
|
|||||||
};
|
};
|
||||||
use aether_data_contracts::repository::candidate_selection::{
|
use aether_data_contracts::repository::candidate_selection::{
|
||||||
StoredMinimalCandidateSelectionRow, StoredPoolKeyCandidateOrder,
|
StoredMinimalCandidateSelectionRow, StoredPoolKeyCandidateOrder,
|
||||||
StoredPoolKeyCandidateRowsQuery,
|
StoredPoolKeyCandidateRowsByKeyIdsQuery, StoredPoolKeyCandidateRowsQuery,
|
||||||
|
};
|
||||||
|
use aether_data_contracts::repository::pool_scores::{
|
||||||
|
ListRankedPoolMembersQuery, PoolMemberHardState, PoolMemberIdentity,
|
||||||
|
PoolMemberScheduleFeedback, PoolScoreScope, StoredPoolMemberScore, POOL_KIND_PROVIDER_KEY_POOL,
|
||||||
};
|
};
|
||||||
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
|
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
|
||||||
use serde_json::{Map, Value};
|
use serde_json::{Map, Value};
|
||||||
@@ -38,6 +42,8 @@ use crate::handlers::shared::{
|
|||||||
use crate::orchestration::LocalExecutionCandidateMetadata;
|
use crate::orchestration::LocalExecutionCandidateMetadata;
|
||||||
use crate::provider_key_auth::provider_key_auth_semantics;
|
use crate::provider_key_auth::provider_key_auth_semantics;
|
||||||
|
|
||||||
|
use super::pool_scores::provider_key_pool_score_scope;
|
||||||
|
|
||||||
static LOAD_BALANCE_SEQUENCE: AtomicU64 = AtomicU64::new(0);
|
static LOAD_BALANCE_SEQUENCE: AtomicU64 = AtomicU64::new(0);
|
||||||
const DEFAULT_POOL_KEY_PAGE_SIZE: u32 = 128;
|
const DEFAULT_POOL_KEY_PAGE_SIZE: u32 = 128;
|
||||||
const DEFAULT_POOL_MAX_SCANNED_KEYS: u32 = 1024;
|
const DEFAULT_POOL_MAX_SCANNED_KEYS: u32 = 1024;
|
||||||
@@ -179,6 +185,8 @@ pub(crate) struct PoolKeyCursor<'a> {
|
|||||||
scanned_keys: u32,
|
scanned_keys: u32,
|
||||||
page_size: u32,
|
page_size: u32,
|
||||||
max_scanned_keys: u32,
|
max_scanned_keys: u32,
|
||||||
|
score_top_n: u32,
|
||||||
|
score_phase_loaded: bool,
|
||||||
skip_reason_counts: BTreeMap<&'static str, u32>,
|
skip_reason_counts: BTreeMap<&'static str, u32>,
|
||||||
next_pool_key_index: u32,
|
next_pool_key_index: u32,
|
||||||
sticky_candidate_loaded: bool,
|
sticky_candidate_loaded: bool,
|
||||||
@@ -199,6 +207,17 @@ impl<'a> PoolKeyCursor<'a> {
|
|||||||
request_auth_channel: Option<&str>,
|
request_auth_channel: Option<&str>,
|
||||||
) -> Self {
|
) -> Self {
|
||||||
let pool_key_order = pool_key_candidate_order_for_group(&group);
|
let pool_key_order = pool_key_candidate_order_for_group(&group);
|
||||||
|
let pool_config = pool_config_for_candidate(&group);
|
||||||
|
let score_top_n = pool_config
|
||||||
|
.as_ref()
|
||||||
|
.map(|config| config.score_top_n)
|
||||||
|
.unwrap_or(u64::from(DEFAULT_POOL_KEY_PAGE_SIZE))
|
||||||
|
.clamp(1, u64::from(u32::MAX)) as u32;
|
||||||
|
let max_scanned_keys = pool_config
|
||||||
|
.as_ref()
|
||||||
|
.map(|config| config.score_fallback_scan_limit)
|
||||||
|
.unwrap_or(u64::from(DEFAULT_POOL_MAX_SCANNED_KEYS))
|
||||||
|
.clamp(1, u64::from(u32::MAX)) as u32;
|
||||||
Self {
|
Self {
|
||||||
state,
|
state,
|
||||||
group,
|
group,
|
||||||
@@ -211,7 +230,9 @@ impl<'a> PoolKeyCursor<'a> {
|
|||||||
next_offset: 0,
|
next_offset: 0,
|
||||||
scanned_keys: 0,
|
scanned_keys: 0,
|
||||||
page_size: DEFAULT_POOL_KEY_PAGE_SIZE,
|
page_size: DEFAULT_POOL_KEY_PAGE_SIZE,
|
||||||
max_scanned_keys: DEFAULT_POOL_MAX_SCANNED_KEYS,
|
max_scanned_keys,
|
||||||
|
score_top_n,
|
||||||
|
score_phase_loaded: false,
|
||||||
skip_reason_counts: BTreeMap::new(),
|
skip_reason_counts: BTreeMap::new(),
|
||||||
next_pool_key_index: 0,
|
next_pool_key_index: 0,
|
||||||
sticky_candidate_loaded: false,
|
sticky_candidate_loaded: false,
|
||||||
@@ -310,6 +331,13 @@ impl<'a> PoolKeyCursor<'a> {
|
|||||||
}
|
}
|
||||||
|
|
||||||
async fn next_page_candidates(&mut self) -> Option<Vec<EligibleLocalExecutionCandidate>> {
|
async fn next_page_candidates(&mut self) -> Option<Vec<EligibleLocalExecutionCandidate>> {
|
||||||
|
if !self.score_phase_loaded {
|
||||||
|
self.score_phase_loaded = true;
|
||||||
|
if let Some(score_candidates) = self.next_score_candidates().await {
|
||||||
|
return Some(score_candidates);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
if self.scanned_keys >= self.max_scanned_keys {
|
if self.scanned_keys >= self.max_scanned_keys {
|
||||||
return None;
|
return None;
|
||||||
}
|
}
|
||||||
@@ -359,6 +387,82 @@ impl<'a> PoolKeyCursor<'a> {
|
|||||||
Some(self.build_page_eligible_candidates(rows).await)
|
Some(self.build_page_eligible_candidates(rows).await)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
async fn next_score_candidates(&mut self) -> Option<Vec<EligibleLocalExecutionCandidate>> {
|
||||||
|
let scope = provider_key_pool_score_scope(
|
||||||
|
self.group.candidate.endpoint_api_format.as_str(),
|
||||||
|
Some(self.group.candidate.model_id.as_str()),
|
||||||
|
);
|
||||||
|
let query = ListRankedPoolMembersQuery {
|
||||||
|
pool_kind: POOL_KIND_PROVIDER_KEY_POOL.to_string(),
|
||||||
|
pool_id: self.group.candidate.provider_id.clone(),
|
||||||
|
capability: scope.capability.clone(),
|
||||||
|
scope_kind: scope.scope_kind.clone(),
|
||||||
|
scope_id: scope.scope_id.clone(),
|
||||||
|
hard_states: vec![PoolMemberHardState::Available, PoolMemberHardState::Unknown],
|
||||||
|
probe_statuses: None,
|
||||||
|
offset: 0,
|
||||||
|
limit: self.score_top_n as usize,
|
||||||
|
};
|
||||||
|
let scores = match self.state.app().data.list_ranked_pool_members(&query).await {
|
||||||
|
Ok(scores) => scores,
|
||||||
|
Err(err) => {
|
||||||
|
warn!(
|
||||||
|
event_name = "pool_group_score_load_failed",
|
||||||
|
log_type = "event",
|
||||||
|
provider_id = %self.group.candidate.provider_id,
|
||||||
|
endpoint_id = %self.group.candidate.endpoint_id,
|
||||||
|
model_id = %self.group.candidate.model_id,
|
||||||
|
selected_provider_model_name = %self.group.candidate.selected_provider_model_name,
|
||||||
|
error = ?err,
|
||||||
|
"gateway pool scheduler failed to read ranked pool member scores"
|
||||||
|
);
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
};
|
||||||
|
if scores.is_empty() {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
|
||||||
|
self.record_score_schedule_interest(&scores).await;
|
||||||
|
|
||||||
|
let key_ids = scores
|
||||||
|
.iter()
|
||||||
|
.map(|score| score.member_id.clone())
|
||||||
|
.collect::<Vec<_>>();
|
||||||
|
let rows_query = StoredPoolKeyCandidateRowsByKeyIdsQuery {
|
||||||
|
api_format: self.group.candidate.endpoint_api_format.clone(),
|
||||||
|
provider_id: self.group.candidate.provider_id.clone(),
|
||||||
|
endpoint_id: self.group.candidate.endpoint_id.clone(),
|
||||||
|
model_id: self.group.candidate.model_id.clone(),
|
||||||
|
selected_provider_model_name: self.group.candidate.selected_provider_model_name.clone(),
|
||||||
|
key_ids,
|
||||||
|
};
|
||||||
|
let rows = match self
|
||||||
|
.state
|
||||||
|
.app()
|
||||||
|
.list_pool_key_candidate_rows_for_group_key_ids(&rows_query)
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
Ok(rows) => rows,
|
||||||
|
Err(err) => {
|
||||||
|
warn!(
|
||||||
|
event_name = "pool_group_score_key_load_failed",
|
||||||
|
log_type = "event",
|
||||||
|
provider_id = %self.group.candidate.provider_id,
|
||||||
|
endpoint_id = %self.group.candidate.endpoint_id,
|
||||||
|
model_id = %self.group.candidate.model_id,
|
||||||
|
selected_provider_model_name = %self.group.candidate.selected_provider_model_name,
|
||||||
|
score_count = scores.len(),
|
||||||
|
error = ?err,
|
||||||
|
"gateway pool scheduler failed to materialize ranked pool keys"
|
||||||
|
);
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
};
|
||||||
|
self.scanned_keys = self.scanned_keys.saturating_add(scores.len() as u32);
|
||||||
|
Some(self.build_page_eligible_candidates(rows).await)
|
||||||
|
}
|
||||||
|
|
||||||
async fn sticky_candidate(&mut self) -> Option<EligibleLocalExecutionCandidate> {
|
async fn sticky_candidate(&mut self) -> Option<EligibleLocalExecutionCandidate> {
|
||||||
let pool_config = pool_config_for_candidate(&self.group)?;
|
let pool_config = pool_config_for_candidate(&self.group)?;
|
||||||
if !admin_provider_pool_cache_affinity_enabled(&pool_config) {
|
if !admin_provider_pool_cache_affinity_enabled(&pool_config) {
|
||||||
@@ -518,6 +622,62 @@ impl<'a> PoolKeyCursor<'a> {
|
|||||||
self.skipped_candidates.append(&mut skipped);
|
self.skipped_candidates.append(&mut skipped);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
async fn record_score_schedule_interest(&self, scores: &[StoredPoolMemberScore]) {
|
||||||
|
if scores.is_empty() {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
let scheduled_at = current_unix_ms() / 1000;
|
||||||
|
let mut failed = 0usize;
|
||||||
|
for score in scores {
|
||||||
|
let identity = PoolMemberIdentity {
|
||||||
|
pool_kind: score.pool_kind.clone(),
|
||||||
|
pool_id: score.pool_id.clone(),
|
||||||
|
member_kind: score.member_kind.clone(),
|
||||||
|
member_id: score.member_id.clone(),
|
||||||
|
};
|
||||||
|
let scope = PoolScoreScope {
|
||||||
|
capability: score.capability.clone(),
|
||||||
|
scope_kind: score.scope_kind.clone(),
|
||||||
|
scope_id: score.scope_id.clone(),
|
||||||
|
};
|
||||||
|
let result = self
|
||||||
|
.state
|
||||||
|
.app()
|
||||||
|
.data
|
||||||
|
.record_pool_member_schedule_feedback(PoolMemberScheduleFeedback {
|
||||||
|
identity,
|
||||||
|
scope: Some(scope),
|
||||||
|
scheduled_at,
|
||||||
|
succeeded: None,
|
||||||
|
hard_state: None,
|
||||||
|
score_delta: None,
|
||||||
|
score_reason_patch: Some(serde_json::json!({
|
||||||
|
"last_schedule_interest": {
|
||||||
|
"provider_id": self.group.candidate.provider_id.as_str(),
|
||||||
|
"endpoint_id": self.group.candidate.endpoint_id.as_str(),
|
||||||
|
"model_id": self.group.candidate.model_id.as_str()
|
||||||
|
}
|
||||||
|
})),
|
||||||
|
})
|
||||||
|
.await;
|
||||||
|
if result.is_err() {
|
||||||
|
failed += 1;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if failed > 0 {
|
||||||
|
warn!(
|
||||||
|
event_name = "pool_group_score_interest_update_failed",
|
||||||
|
log_type = "event",
|
||||||
|
provider_id = %self.group.candidate.provider_id,
|
||||||
|
endpoint_id = %self.group.candidate.endpoint_id,
|
||||||
|
model_id = %self.group.candidate.model_id,
|
||||||
|
failed_count = failed,
|
||||||
|
score_count = scores.len(),
|
||||||
|
"gateway pool scheduler failed to record some pool score schedule interests"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
async fn build_page_eligible_candidates(
|
async fn build_page_eligible_candidates(
|
||||||
&mut self,
|
&mut self,
|
||||||
rows: Vec<StoredMinimalCandidateSelectionRow>,
|
rows: Vec<StoredMinimalCandidateSelectionRow>,
|
||||||
|
|||||||
196
apps/aether-gateway/src/ai_serving/planner/pool_scores.rs
Normal file
196
apps/aether-gateway/src/ai_serving/planner/pool_scores.rs
Normal file
@@ -0,0 +1,196 @@
|
|||||||
|
use aether_ai_serving::{
|
||||||
|
score_pool_member_with_rules, PoolMemberScoreInput, PoolMemberScoreRules, POOL_SCORE_VERSION,
|
||||||
|
};
|
||||||
|
use aether_data_contracts::repository::pool_scores::{
|
||||||
|
PoolMemberIdentity, PoolMemberProbeStatus, PoolScoreScope, UpsertPoolMemberScore,
|
||||||
|
POOL_SCORE_SCOPE_KIND_MODEL,
|
||||||
|
};
|
||||||
|
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
|
||||||
|
use serde_json::Value;
|
||||||
|
|
||||||
|
use crate::handlers::shared::{provider_key_health_summary, provider_key_status_snapshot_payload};
|
||||||
|
|
||||||
|
pub(crate) fn build_provider_key_pool_score_upsert(
|
||||||
|
key: &StoredProviderCatalogKey,
|
||||||
|
provider_type: &str,
|
||||||
|
api_format: &str,
|
||||||
|
model_id: Option<&str>,
|
||||||
|
existing: Option<&aether_data_contracts::repository::pool_scores::StoredPoolMemberScore>,
|
||||||
|
now_unix_secs: u64,
|
||||||
|
score_rules: PoolMemberScoreRules,
|
||||||
|
) -> UpsertPoolMemberScore {
|
||||||
|
let identity = PoolMemberIdentity::provider_api_key(key.provider_id.clone(), key.id.clone());
|
||||||
|
let scope = provider_key_pool_score_scope(api_format, model_id);
|
||||||
|
let input = provider_key_score_input(
|
||||||
|
key,
|
||||||
|
provider_type,
|
||||||
|
identity.clone(),
|
||||||
|
scope.clone(),
|
||||||
|
existing,
|
||||||
|
now_unix_secs,
|
||||||
|
);
|
||||||
|
let output = score_pool_member_with_rules(&input, score_rules);
|
||||||
|
UpsertPoolMemberScore {
|
||||||
|
id: provider_key_pool_score_id(&identity, &scope),
|
||||||
|
identity,
|
||||||
|
scope,
|
||||||
|
score: output.score,
|
||||||
|
hard_state: output.hard_state,
|
||||||
|
score_version: POOL_SCORE_VERSION,
|
||||||
|
score_reason: output.score_reason,
|
||||||
|
last_ranked_at: Some(now_unix_secs),
|
||||||
|
last_scheduled_at: existing.and_then(|score| score.last_scheduled_at),
|
||||||
|
last_success_at: existing.and_then(|score| score.last_success_at),
|
||||||
|
last_failure_at: existing.and_then(|score| score.last_failure_at),
|
||||||
|
failure_count: existing.map(|score| score.failure_count).unwrap_or(0),
|
||||||
|
last_probe_attempt_at: existing.and_then(|score| score.last_probe_attempt_at),
|
||||||
|
last_probe_success_at: existing.and_then(|score| score.last_probe_success_at),
|
||||||
|
last_probe_failure_at: existing.and_then(|score| score.last_probe_failure_at),
|
||||||
|
probe_failure_count: existing.map(|score| score.probe_failure_count).unwrap_or(0),
|
||||||
|
probe_status: existing
|
||||||
|
.map(|score| score.probe_status)
|
||||||
|
.unwrap_or(PoolMemberProbeStatus::Never),
|
||||||
|
updated_at: now_unix_secs,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(crate) fn provider_key_pool_score_scope(
|
||||||
|
api_format: &str,
|
||||||
|
model_id: Option<&str>,
|
||||||
|
) -> PoolScoreScope {
|
||||||
|
PoolScoreScope {
|
||||||
|
capability: api_format.trim().to_ascii_lowercase(),
|
||||||
|
scope_kind: POOL_SCORE_SCOPE_KIND_MODEL.to_string(),
|
||||||
|
scope_id: model_id
|
||||||
|
.map(str::trim)
|
||||||
|
.filter(|value| !value.is_empty())
|
||||||
|
.map(ToOwned::to_owned),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(crate) fn provider_key_pool_score_id(
|
||||||
|
identity: &PoolMemberIdentity,
|
||||||
|
scope: &PoolScoreScope,
|
||||||
|
) -> String {
|
||||||
|
let raw = format!(
|
||||||
|
"{}:{}:{}:{}:{}:{}:{}",
|
||||||
|
identity.pool_kind,
|
||||||
|
identity.pool_id,
|
||||||
|
identity.member_kind,
|
||||||
|
identity.member_id,
|
||||||
|
scope.capability,
|
||||||
|
scope.scope_kind,
|
||||||
|
scope.scope_id.as_deref().unwrap_or("*")
|
||||||
|
);
|
||||||
|
format!(
|
||||||
|
"pms-{:016x}-{:016x}",
|
||||||
|
stable_hash(raw.as_bytes()),
|
||||||
|
stable_hash(identity.member_id.as_bytes())
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn provider_key_score_input(
|
||||||
|
key: &StoredProviderCatalogKey,
|
||||||
|
provider_type: &str,
|
||||||
|
identity: PoolMemberIdentity,
|
||||||
|
scope: PoolScoreScope,
|
||||||
|
existing: Option<&aether_data_contracts::repository::pool_scores::StoredPoolMemberScore>,
|
||||||
|
now_unix_secs: u64,
|
||||||
|
) -> PoolMemberScoreInput {
|
||||||
|
let status_snapshot = provider_key_status_snapshot_payload(key, provider_type);
|
||||||
|
let quota_snapshot = status_snapshot
|
||||||
|
.as_object()
|
||||||
|
.and_then(|snapshot| snapshot.get("quota"))
|
||||||
|
.and_then(Value::as_object);
|
||||||
|
let account_snapshot = status_snapshot
|
||||||
|
.as_object()
|
||||||
|
.and_then(|snapshot| snapshot.get("account"))
|
||||||
|
.and_then(Value::as_object);
|
||||||
|
let (health_score, _, _, any_circuit_open, _) = provider_key_health_summary(key);
|
||||||
|
let health_score = key
|
||||||
|
.health_by_format
|
||||||
|
.as_ref()
|
||||||
|
.and_then(Value::as_object)
|
||||||
|
.filter(|payload| !payload.is_empty())
|
||||||
|
.map(|_| health_score);
|
||||||
|
|
||||||
|
PoolMemberScoreInput {
|
||||||
|
identity,
|
||||||
|
scope: scope.clone(),
|
||||||
|
internal_priority: key.internal_priority,
|
||||||
|
is_active: key.is_active,
|
||||||
|
health_score,
|
||||||
|
quota_usage_ratio: quota_snapshot
|
||||||
|
.and_then(|quota| quota.get("usage_ratio"))
|
||||||
|
.and_then(json_f64)
|
||||||
|
.map(|value| value.clamp(0.0, 1.0)),
|
||||||
|
quota_exhausted: quota_snapshot
|
||||||
|
.and_then(|quota| quota.get("exhausted"))
|
||||||
|
.and_then(Value::as_bool)
|
||||||
|
.unwrap_or(false),
|
||||||
|
account_blocked: account_snapshot
|
||||||
|
.and_then(|account| account.get("blocked"))
|
||||||
|
.and_then(Value::as_bool)
|
||||||
|
.unwrap_or(false),
|
||||||
|
oauth_invalid_reason: key.oauth_invalid_reason.clone(),
|
||||||
|
circuit_open: any_circuit_open
|
||||||
|
|| circuit_open_for_scope(key.circuit_breaker_by_format.as_ref(), &scope),
|
||||||
|
success_count: key.success_count.unwrap_or(0).into(),
|
||||||
|
error_count: key.error_count.unwrap_or(0).into(),
|
||||||
|
total_response_time_ms: key.total_response_time_ms.unwrap_or(0).into(),
|
||||||
|
total_tokens: key.total_tokens,
|
||||||
|
total_cost_usd: key.total_cost_usd,
|
||||||
|
last_used_at: key.last_used_at_unix_secs,
|
||||||
|
last_probe_success_at: existing.and_then(|score| score.last_probe_success_at),
|
||||||
|
probe_failure_count: existing.map(|score| score.probe_failure_count).unwrap_or(0),
|
||||||
|
probe_status: existing
|
||||||
|
.map(|score| score.probe_status)
|
||||||
|
.unwrap_or(PoolMemberProbeStatus::Never),
|
||||||
|
now_unix_secs,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn circuit_open_for_scope(circuit_by_format: Option<&Value>, scope: &PoolScoreScope) -> bool {
|
||||||
|
let Some(formats) = circuit_by_format.and_then(Value::as_object) else {
|
||||||
|
return false;
|
||||||
|
};
|
||||||
|
let keys = api_format_lookup_keys(&scope.capability);
|
||||||
|
keys.iter().any(|key| {
|
||||||
|
formats
|
||||||
|
.get(key)
|
||||||
|
.and_then(Value::as_object)
|
||||||
|
.and_then(|value| value.get("open"))
|
||||||
|
.and_then(Value::as_bool)
|
||||||
|
.unwrap_or(false)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn json_f64(value: &Value) -> Option<f64> {
|
||||||
|
value.as_f64().or_else(|| {
|
||||||
|
value
|
||||||
|
.as_str()
|
||||||
|
.map(str::trim)
|
||||||
|
.filter(|value| !value.is_empty())
|
||||||
|
.and_then(|value| value.parse::<f64>().ok())
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn api_format_lookup_keys(api_format: &str) -> Vec<String> {
|
||||||
|
let normalized = crate::ai_serving::normalize_api_format_alias(api_format);
|
||||||
|
let mut keys = crate::ai_serving::api_format_storage_aliases(&normalized);
|
||||||
|
if !keys.iter().any(|value| value == &normalized) {
|
||||||
|
keys.push(normalized);
|
||||||
|
}
|
||||||
|
keys.sort();
|
||||||
|
keys.dedup();
|
||||||
|
keys
|
||||||
|
}
|
||||||
|
|
||||||
|
fn stable_hash(bytes: &[u8]) -> u64 {
|
||||||
|
let mut hash = 0xcbf29ce484222325u64;
|
||||||
|
for byte in bytes {
|
||||||
|
hash ^= u64::from(*byte);
|
||||||
|
hash = hash.wrapping_mul(0x100000001b3);
|
||||||
|
}
|
||||||
|
hash
|
||||||
|
}
|
||||||
@@ -2,9 +2,10 @@ pub(crate) use aether_ai_formats::api::{
|
|||||||
aggregate_claude_stream_sync_response, aggregate_gemini_stream_sync_response,
|
aggregate_claude_stream_sync_response, aggregate_gemini_stream_sync_response,
|
||||||
aggregate_openai_chat_stream_sync_response, aggregate_openai_responses_stream_sync_response,
|
aggregate_openai_chat_stream_sync_response, aggregate_openai_responses_stream_sync_response,
|
||||||
aggregate_standard_chat_stream_sync_response, aggregate_standard_cli_stream_sync_response,
|
aggregate_standard_chat_stream_sync_response, aggregate_standard_cli_stream_sync_response,
|
||||||
api_format_alias_matches, apply_codex_openai_responses_special_body_edits,
|
api_format_alias_matches, api_format_storage_aliases,
|
||||||
apply_codex_openai_responses_special_headers, apply_model_directive_mapping_patch,
|
apply_codex_openai_responses_special_body_edits, apply_codex_openai_responses_special_headers,
|
||||||
apply_model_directive_overrides_from_model, apply_model_directive_overrides_from_request,
|
apply_model_directive_mapping_patch, apply_model_directive_overrides_from_model,
|
||||||
|
apply_model_directive_overrides_from_request,
|
||||||
apply_openai_responses_compact_special_body_edits, build_chatgpt_web_image_request_body,
|
apply_openai_responses_compact_special_body_edits, build_chatgpt_web_image_request_body,
|
||||||
build_core_error_body_for_client_format, build_cross_format_openai_chat_request_body,
|
build_core_error_body_for_client_format, build_cross_format_openai_chat_request_body,
|
||||||
build_cross_format_openai_chat_request_body_with_model_directives,
|
build_cross_format_openai_chat_request_body_with_model_directives,
|
||||||
|
|||||||
@@ -240,6 +240,18 @@ pub(super) fn classify_admin_observability_family_route(
|
|||||||
"admin:pool",
|
"admin:pool",
|
||||||
false,
|
false,
|
||||||
))
|
))
|
||||||
|
} else if method == http::Method::GET
|
||||||
|
&& normalized_path_no_trailing.starts_with("/api/admin/pool/")
|
||||||
|
&& normalized_path_no_trailing.ends_with("/scores")
|
||||||
|
&& normalized_path_no_trailing.matches('/').count() == 5
|
||||||
|
{
|
||||||
|
Some(classified(
|
||||||
|
"admin_proxy",
|
||||||
|
"pool_manage",
|
||||||
|
"scores",
|
||||||
|
"admin:pool",
|
||||||
|
false,
|
||||||
|
))
|
||||||
} else if method == http::Method::POST
|
} else if method == http::Method::POST
|
||||||
&& normalized_path_no_trailing.starts_with("/api/admin/pool/")
|
&& normalized_path_no_trailing.starts_with("/api/admin/pool/")
|
||||||
&& normalized_path_no_trailing.ends_with("/keys/batch-import")
|
&& normalized_path_no_trailing.ends_with("/keys/batch-import")
|
||||||
|
|||||||
@@ -52,6 +52,14 @@ fn classifies_admin_pool_provider_key_routes_as_admin_proxy_route() {
|
|||||||
assert_eq!(list.route_family.as_deref(), Some("pool_manage"));
|
assert_eq!(list.route_family.as_deref(), Some("pool_manage"));
|
||||||
assert_eq!(list.route_kind.as_deref(), Some("list_keys"));
|
assert_eq!(list.route_kind.as_deref(), Some("list_keys"));
|
||||||
|
|
||||||
|
let scores_uri: Uri = "/api/admin/pool/provider-1/scores?api_format=openai:responses"
|
||||||
|
.parse()
|
||||||
|
.expect("uri should parse");
|
||||||
|
let scores = classify_control_route(&http::Method::GET, &scores_uri, &headers)
|
||||||
|
.expect("route should classify");
|
||||||
|
assert_eq!(scores.route_family.as_deref(), Some("pool_manage"));
|
||||||
|
assert_eq!(scores.route_kind.as_deref(), Some("scores"));
|
||||||
|
|
||||||
let batch_import_uri: Uri = "/api/admin/pool/provider-1/keys/batch-import"
|
let batch_import_uri: Uri = "/api/admin/pool/provider-1/keys/batch-import"
|
||||||
.parse()
|
.parse()
|
||||||
.expect("uri should parse");
|
.expect("uri should parse");
|
||||||
|
|||||||
@@ -45,6 +45,8 @@ impl GatewayDataState {
|
|||||||
request_candidate_writer: None,
|
request_candidate_writer: None,
|
||||||
provider_catalog_reader: None,
|
provider_catalog_reader: None,
|
||||||
provider_catalog_writer: None,
|
provider_catalog_writer: None,
|
||||||
|
pool_score_reader: None,
|
||||||
|
pool_score_writer: None,
|
||||||
provider_quota_reader: None,
|
provider_quota_reader: None,
|
||||||
provider_quota_writer: None,
|
provider_quota_writer: None,
|
||||||
usage_reader: None,
|
usage_reader: None,
|
||||||
@@ -86,6 +88,8 @@ impl GatewayDataState {
|
|||||||
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();
|
||||||
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_writer = backends.write().pool_scores();
|
||||||
let provider_quota_reader = backends.read().provider_quotas();
|
let provider_quota_reader = backends.read().provider_quotas();
|
||||||
let provider_quota_writer = backends.write().provider_quotas();
|
let provider_quota_writer = backends.write().provider_quotas();
|
||||||
let usage_reader = backends.read().usage();
|
let usage_reader = backends.read().usage();
|
||||||
@@ -125,6 +129,8 @@ impl GatewayDataState {
|
|||||||
request_candidate_writer,
|
request_candidate_writer,
|
||||||
provider_catalog_reader,
|
provider_catalog_reader,
|
||||||
provider_catalog_writer,
|
provider_catalog_writer,
|
||||||
|
pool_score_reader,
|
||||||
|
pool_score_writer,
|
||||||
provider_quota_reader,
|
provider_quota_reader,
|
||||||
provider_quota_writer,
|
provider_quota_writer,
|
||||||
usage_reader,
|
usage_reader,
|
||||||
@@ -263,6 +269,14 @@ impl GatewayDataState {
|
|||||||
self.provider_catalog_writer.is_some()
|
self.provider_catalog_writer.is_some()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub(crate) fn has_pool_score_reader(&self) -> bool {
|
||||||
|
self.pool_score_reader.is_some()
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(crate) fn has_pool_score_writer(&self) -> bool {
|
||||||
|
self.pool_score_writer.is_some()
|
||||||
|
}
|
||||||
|
|
||||||
pub(crate) fn has_proxy_node_reader(&self) -> bool {
|
pub(crate) fn has_proxy_node_reader(&self) -> bool {
|
||||||
self.proxy_node_reader.is_some()
|
self.proxy_node_reader.is_some()
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -93,7 +93,8 @@ use aether_data_contracts::repository::billing::{
|
|||||||
};
|
};
|
||||||
use aether_data_contracts::repository::candidate_selection::{
|
use aether_data_contracts::repository::candidate_selection::{
|
||||||
MinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
|
MinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
|
||||||
StoredPoolKeyCandidateRowsQuery, StoredRequestedModelCandidateRowsQuery,
|
StoredPoolKeyCandidateRowsByKeyIdsQuery, StoredPoolKeyCandidateRowsQuery,
|
||||||
|
StoredRequestedModelCandidateRowsQuery,
|
||||||
};
|
};
|
||||||
use aether_data_contracts::repository::candidates::{
|
use aether_data_contracts::repository::candidates::{
|
||||||
PublicHealthStatusCount, PublicHealthTimelineBucket, RequestCandidateReadRepository,
|
PublicHealthStatusCount, PublicHealthTimelineBucket, RequestCandidateReadRepository,
|
||||||
@@ -107,6 +108,13 @@ use aether_data_contracts::repository::global_models::{
|
|||||||
StoredProviderModelStats, StoredPublicCatalogModel, StoredPublicGlobalModel,
|
StoredProviderModelStats, StoredPublicCatalogModel, StoredPublicGlobalModel,
|
||||||
StoredPublicGlobalModelPage, UpdateAdminGlobalModelRecord, UpsertAdminProviderModelRecord,
|
StoredPublicGlobalModelPage, UpdateAdminGlobalModelRecord, UpsertAdminProviderModelRecord,
|
||||||
};
|
};
|
||||||
|
use aether_data_contracts::repository::pool_scores::{
|
||||||
|
GetPoolMemberScoresByIdsQuery, ListPoolMemberProbeCandidatesQuery, ListPoolMemberScoresQuery,
|
||||||
|
ListRankedPoolMembersQuery, PoolMemberHardState, PoolMemberIdentity, PoolMemberProbeAttempt,
|
||||||
|
PoolMemberProbeResult, PoolMemberProbeStatus, PoolMemberScheduleFeedback,
|
||||||
|
PoolMemberScoreWriteRepository, PoolScoreReadRepository, PoolScoreScope, StoredPoolMemberScore,
|
||||||
|
UpsertPoolMemberScore,
|
||||||
|
};
|
||||||
use aether_data_contracts::repository::provider_catalog::{
|
use aether_data_contracts::repository::provider_catalog::{
|
||||||
ProviderCatalogKeyListQuery, ProviderCatalogReadRepository, ProviderCatalogWriteRepository,
|
ProviderCatalogKeyListQuery, ProviderCatalogReadRepository, ProviderCatalogWriteRepository,
|
||||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogKeyPage,
|
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogKeyPage,
|
||||||
@@ -156,6 +164,8 @@ pub(crate) struct GatewayDataState {
|
|||||||
request_candidate_writer: Option<Arc<dyn RequestCandidateWriteRepository>>,
|
request_candidate_writer: Option<Arc<dyn RequestCandidateWriteRepository>>,
|
||||||
provider_catalog_reader: Option<Arc<dyn ProviderCatalogReadRepository>>,
|
provider_catalog_reader: Option<Arc<dyn ProviderCatalogReadRepository>>,
|
||||||
provider_catalog_writer: Option<Arc<dyn ProviderCatalogWriteRepository>>,
|
provider_catalog_writer: Option<Arc<dyn ProviderCatalogWriteRepository>>,
|
||||||
|
pool_score_reader: Option<Arc<dyn PoolScoreReadRepository>>,
|
||||||
|
pool_score_writer: Option<Arc<dyn PoolMemberScoreWriteRepository>>,
|
||||||
provider_quota_reader: Option<Arc<dyn ProviderQuotaReadRepository>>,
|
provider_quota_reader: Option<Arc<dyn ProviderQuotaReadRepository>>,
|
||||||
provider_quota_writer: Option<Arc<dyn ProviderQuotaWriteRepository>>,
|
provider_quota_writer: Option<Arc<dyn ProviderQuotaWriteRepository>>,
|
||||||
usage_reader: Option<Arc<dyn UsageReadRepository>>,
|
usage_reader: Option<Arc<dyn UsageReadRepository>>,
|
||||||
@@ -257,6 +267,8 @@ impl fmt::Debug for GatewayDataState {
|
|||||||
"has_provider_catalog_writer",
|
"has_provider_catalog_writer",
|
||||||
&self.provider_catalog_writer.is_some(),
|
&self.provider_catalog_writer.is_some(),
|
||||||
)
|
)
|
||||||
|
.field("has_pool_score_reader", &self.pool_score_reader.is_some())
|
||||||
|
.field("has_pool_score_writer", &self.pool_score_writer.is_some())
|
||||||
.field(
|
.field(
|
||||||
"has_provider_quota_reader",
|
"has_provider_quota_reader",
|
||||||
&self.provider_quota_reader.is_some(),
|
&self.provider_quota_reader.is_some(),
|
||||||
@@ -287,6 +299,7 @@ mod catalog;
|
|||||||
mod core;
|
mod core;
|
||||||
mod integrations;
|
mod integrations;
|
||||||
mod models;
|
mod models;
|
||||||
|
mod pool_scores;
|
||||||
mod runtime;
|
mod runtime;
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod testing;
|
mod testing;
|
||||||
|
|||||||
@@ -2,7 +2,8 @@ use super::{
|
|||||||
AdminGlobalModelListQuery, AdminProviderModelListQuery, CreateAdminGlobalModelRecord,
|
AdminGlobalModelListQuery, AdminProviderModelListQuery, CreateAdminGlobalModelRecord,
|
||||||
DataLayerError, GatewayDataState, PublicCatalogModelListQuery, PublicCatalogModelSearchQuery,
|
DataLayerError, GatewayDataState, PublicCatalogModelListQuery, PublicCatalogModelSearchQuery,
|
||||||
PublicGlobalModelQuery, StoredAdminGlobalModel, StoredAdminGlobalModelPage,
|
PublicGlobalModelQuery, StoredAdminGlobalModel, StoredAdminGlobalModelPage,
|
||||||
StoredAdminProviderModel, StoredMinimalCandidateSelectionRow, StoredPoolKeyCandidateRowsQuery,
|
StoredAdminProviderModel, StoredMinimalCandidateSelectionRow,
|
||||||
|
StoredPoolKeyCandidateRowsByKeyIdsQuery, StoredPoolKeyCandidateRowsQuery,
|
||||||
StoredProviderActiveGlobalModel, StoredProviderModelStats, StoredPublicCatalogModel,
|
StoredProviderActiveGlobalModel, StoredProviderModelStats, StoredPublicCatalogModel,
|
||||||
StoredPublicGlobalModel, StoredPublicGlobalModelPage, StoredRequestedModelCandidateRowsQuery,
|
StoredPublicGlobalModel, StoredPublicGlobalModelPage, StoredRequestedModelCandidateRowsQuery,
|
||||||
UpdateAdminGlobalModelRecord, UpsertAdminProviderModelRecord,
|
UpdateAdminGlobalModelRecord, UpsertAdminProviderModelRecord,
|
||||||
@@ -73,6 +74,16 @@ impl GatewayDataState {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub(crate) async fn list_pool_key_candidate_rows_for_group_key_ids(
|
||||||
|
&self,
|
||||||
|
query: &StoredPoolKeyCandidateRowsByKeyIdsQuery,
|
||||||
|
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
|
||||||
|
match &self.minimal_candidate_selection_reader {
|
||||||
|
Some(repository) => repository.list_pool_key_rows_for_group_key_ids(query).await,
|
||||||
|
None => Ok(Vec::new()),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
pub(crate) async fn list_public_global_models(
|
pub(crate) async fn list_public_global_models(
|
||||||
&self,
|
&self,
|
||||||
query: &PublicGlobalModelQuery,
|
query: &PublicGlobalModelQuery,
|
||||||
|
|||||||
123
apps/aether-gateway/src/data/state/pool_scores.rs
Normal file
123
apps/aether-gateway/src/data/state/pool_scores.rs
Normal file
@@ -0,0 +1,123 @@
|
|||||||
|
use super::{
|
||||||
|
DataLayerError, GatewayDataState, GetPoolMemberScoresByIdsQuery,
|
||||||
|
ListPoolMemberProbeCandidatesQuery, ListPoolMemberScoresQuery, ListRankedPoolMembersQuery,
|
||||||
|
PoolMemberHardState, PoolMemberIdentity, PoolMemberProbeAttempt, PoolMemberProbeResult,
|
||||||
|
PoolMemberScheduleFeedback, PoolScoreScope, StoredPoolMemberScore, UpsertPoolMemberScore,
|
||||||
|
};
|
||||||
|
|
||||||
|
impl GatewayDataState {
|
||||||
|
pub(crate) async fn list_ranked_pool_members(
|
||||||
|
&self,
|
||||||
|
query: &ListRankedPoolMembersQuery,
|
||||||
|
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
|
||||||
|
match &self.pool_score_reader {
|
||||||
|
Some(repository) => repository.list_ranked_pool_members(query).await,
|
||||||
|
None => Ok(Vec::new()),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(crate) async fn list_pool_member_probe_candidates(
|
||||||
|
&self,
|
||||||
|
query: &ListPoolMemberProbeCandidatesQuery,
|
||||||
|
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
|
||||||
|
match &self.pool_score_reader {
|
||||||
|
Some(repository) => repository.list_pool_member_probe_candidates(query).await,
|
||||||
|
None => Ok(Vec::new()),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(crate) async fn list_pool_member_scores(
|
||||||
|
&self,
|
||||||
|
query: &ListPoolMemberScoresQuery,
|
||||||
|
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
|
||||||
|
match &self.pool_score_reader {
|
||||||
|
Some(repository) => repository.list_pool_member_scores(query).await,
|
||||||
|
None => Ok(Vec::new()),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(crate) async fn get_pool_member_scores_by_ids(
|
||||||
|
&self,
|
||||||
|
query: &GetPoolMemberScoresByIdsQuery,
|
||||||
|
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
|
||||||
|
match &self.pool_score_reader {
|
||||||
|
Some(repository) => repository.get_pool_member_scores_by_ids(query).await,
|
||||||
|
None => Ok(Vec::new()),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(crate) async fn upsert_pool_member_score(
|
||||||
|
&self,
|
||||||
|
score: UpsertPoolMemberScore,
|
||||||
|
) -> Result<Option<StoredPoolMemberScore>, DataLayerError> {
|
||||||
|
match &self.pool_score_writer {
|
||||||
|
Some(repository) => repository.upsert_pool_member_score(score).await.map(Some),
|
||||||
|
None => Ok(None),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(crate) async fn record_pool_member_probe_result(
|
||||||
|
&self,
|
||||||
|
result: PoolMemberProbeResult,
|
||||||
|
) -> Result<usize, DataLayerError> {
|
||||||
|
match &self.pool_score_writer {
|
||||||
|
Some(repository) => repository.record_pool_member_probe_result(result).await,
|
||||||
|
None => Ok(0),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(crate) async fn mark_pool_member_probe_in_progress(
|
||||||
|
&self,
|
||||||
|
attempt: PoolMemberProbeAttempt,
|
||||||
|
) -> Result<usize, DataLayerError> {
|
||||||
|
match &self.pool_score_writer {
|
||||||
|
Some(repository) => repository.mark_pool_member_probe_in_progress(attempt).await,
|
||||||
|
None => Ok(0),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(crate) async fn record_pool_member_schedule_feedback(
|
||||||
|
&self,
|
||||||
|
feedback: PoolMemberScheduleFeedback,
|
||||||
|
) -> Result<usize, DataLayerError> {
|
||||||
|
match &self.pool_score_writer {
|
||||||
|
Some(repository) => {
|
||||||
|
repository
|
||||||
|
.record_pool_member_schedule_feedback(feedback)
|
||||||
|
.await
|
||||||
|
}
|
||||||
|
None => Ok(0),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(crate) async fn mark_pool_member_hard_state(
|
||||||
|
&self,
|
||||||
|
identity: &PoolMemberIdentity,
|
||||||
|
scope: Option<&PoolScoreScope>,
|
||||||
|
hard_state: PoolMemberHardState,
|
||||||
|
updated_at: u64,
|
||||||
|
) -> Result<usize, DataLayerError> {
|
||||||
|
match &self.pool_score_writer {
|
||||||
|
Some(repository) => {
|
||||||
|
repository
|
||||||
|
.mark_pool_member_hard_state(identity, scope, hard_state, updated_at)
|
||||||
|
.await
|
||||||
|
}
|
||||||
|
None => Ok(0),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(crate) async fn delete_pool_member_scores_for_member(
|
||||||
|
&self,
|
||||||
|
identity: &PoolMemberIdentity,
|
||||||
|
) -> Result<usize, DataLayerError> {
|
||||||
|
match &self.pool_score_writer {
|
||||||
|
Some(repository) => {
|
||||||
|
repository
|
||||||
|
.delete_pool_member_scores_for_member(identity)
|
||||||
|
.await
|
||||||
|
}
|
||||||
|
None => Ok(0),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -34,6 +34,8 @@ impl GatewayDataState {
|
|||||||
request_candidate_writer: None,
|
request_candidate_writer: None,
|
||||||
provider_catalog_reader: None,
|
provider_catalog_reader: None,
|
||||||
provider_catalog_writer: None,
|
provider_catalog_writer: None,
|
||||||
|
pool_score_reader: None,
|
||||||
|
pool_score_writer: None,
|
||||||
provider_quota_reader: None,
|
provider_quota_reader: None,
|
||||||
provider_quota_writer: None,
|
provider_quota_writer: None,
|
||||||
usage_reader: None,
|
usage_reader: None,
|
||||||
@@ -86,6 +88,8 @@ impl GatewayDataState {
|
|||||||
request_candidate_writer: None,
|
request_candidate_writer: None,
|
||||||
provider_catalog_reader: None,
|
provider_catalog_reader: None,
|
||||||
provider_catalog_writer: None,
|
provider_catalog_writer: None,
|
||||||
|
pool_score_reader: None,
|
||||||
|
pool_score_writer: None,
|
||||||
provider_quota_reader: None,
|
provider_quota_reader: None,
|
||||||
provider_quota_writer: None,
|
provider_quota_writer: None,
|
||||||
usage_reader: None,
|
usage_reader: None,
|
||||||
|
|||||||
@@ -2,6 +2,7 @@ use std::collections::BTreeMap;
|
|||||||
use std::sync::{Arc, RwLock};
|
use std::sync::{Arc, RwLock};
|
||||||
|
|
||||||
use aether_data_contracts::repository::candidates::RequestCandidateRepository;
|
use aether_data_contracts::repository::candidates::RequestCandidateRepository;
|
||||||
|
use aether_data_contracts::repository::pool_scores::PoolMemberScoreRepository;
|
||||||
use aether_data_contracts::repository::quota::ProviderQuotaRepository;
|
use aether_data_contracts::repository::quota::ProviderQuotaRepository;
|
||||||
use aether_data_contracts::repository::usage::UsageRepository;
|
use aether_data_contracts::repository::usage::UsageRepository;
|
||||||
|
|
||||||
@@ -12,12 +13,13 @@ use super::{
|
|||||||
GeminiFileMappingWriteRepository, GlobalModelReadRepository, GlobalModelWriteRepository,
|
GeminiFileMappingWriteRepository, GlobalModelReadRepository, GlobalModelWriteRepository,
|
||||||
ManagementTokenReadRepository, ManagementTokenWriteRepository,
|
ManagementTokenReadRepository, ManagementTokenWriteRepository,
|
||||||
MinimalCandidateSelectionReadRepository, OAuthProviderReadRepository,
|
MinimalCandidateSelectionReadRepository, OAuthProviderReadRepository,
|
||||||
OAuthProviderWriteRepository, ProviderCatalogReadRepository, ProviderCatalogWriteRepository,
|
OAuthProviderWriteRepository, PoolMemberScoreWriteRepository, PoolScoreReadRepository,
|
||||||
ProviderQuotaReadRepository, ProviderQuotaWriteRepository, ProxyNodeReadRepository,
|
ProviderCatalogReadRepository, ProviderCatalogWriteRepository, ProviderQuotaReadRepository,
|
||||||
ProxyNodeWriteRepository, RequestCandidateReadRepository, RequestCandidateWriteRepository,
|
ProviderQuotaWriteRepository, ProxyNodeReadRepository, ProxyNodeWriteRepository,
|
||||||
SettlementWriteRepository, StoredSystemConfigEntry, StoredUserPreferenceRecord,
|
RequestCandidateReadRepository, RequestCandidateWriteRepository, SettlementWriteRepository,
|
||||||
UsageReadRepository, UsageWriteRepository, UserReadRepository, VideoTaskReadRepository,
|
StoredSystemConfigEntry, StoredUserPreferenceRecord, UsageReadRepository, UsageWriteRepository,
|
||||||
VideoTaskWriteRepository, WalletReadRepository, WalletWriteRepository,
|
UserReadRepository, VideoTaskReadRepository, VideoTaskWriteRepository, WalletReadRepository,
|
||||||
|
WalletWriteRepository,
|
||||||
};
|
};
|
||||||
|
|
||||||
mod announcements;
|
mod announcements;
|
||||||
@@ -67,6 +69,8 @@ impl GatewayDataState {
|
|||||||
request_candidate_writer: None,
|
request_candidate_writer: None,
|
||||||
provider_catalog_reader: None,
|
provider_catalog_reader: None,
|
||||||
provider_catalog_writer: None,
|
provider_catalog_writer: None,
|
||||||
|
pool_score_reader: None,
|
||||||
|
pool_score_writer: None,
|
||||||
provider_quota_reader: None,
|
provider_quota_reader: None,
|
||||||
provider_quota_writer: None,
|
provider_quota_writer: None,
|
||||||
usage_reader: None,
|
usage_reader: None,
|
||||||
@@ -118,6 +122,8 @@ impl GatewayDataState {
|
|||||||
request_candidate_writer: Some(request_candidate_writer),
|
request_candidate_writer: Some(request_candidate_writer),
|
||||||
provider_catalog_reader: None,
|
provider_catalog_reader: None,
|
||||||
provider_catalog_writer: None,
|
provider_catalog_writer: None,
|
||||||
|
pool_score_reader: None,
|
||||||
|
pool_score_writer: None,
|
||||||
provider_quota_reader: None,
|
provider_quota_reader: None,
|
||||||
provider_quota_writer: None,
|
provider_quota_writer: None,
|
||||||
usage_reader: None,
|
usage_reader: None,
|
||||||
@@ -165,6 +171,8 @@ impl GatewayDataState {
|
|||||||
request_candidate_writer: None,
|
request_candidate_writer: None,
|
||||||
provider_catalog_reader: Some(repository),
|
provider_catalog_reader: Some(repository),
|
||||||
provider_catalog_writer: None,
|
provider_catalog_writer: None,
|
||||||
|
pool_score_reader: None,
|
||||||
|
pool_score_writer: None,
|
||||||
provider_quota_reader: None,
|
provider_quota_reader: None,
|
||||||
provider_quota_writer: None,
|
provider_quota_writer: None,
|
||||||
usage_reader: None,
|
usage_reader: None,
|
||||||
@@ -300,6 +308,8 @@ impl GatewayDataState {
|
|||||||
request_candidate_writer: None,
|
request_candidate_writer: None,
|
||||||
provider_catalog_reader: Some(provider_catalog_reader),
|
provider_catalog_reader: Some(provider_catalog_reader),
|
||||||
provider_catalog_writer: Some(provider_catalog_writer),
|
provider_catalog_writer: Some(provider_catalog_writer),
|
||||||
|
pool_score_reader: None,
|
||||||
|
pool_score_writer: None,
|
||||||
provider_quota_reader: None,
|
provider_quota_reader: None,
|
||||||
provider_quota_writer: None,
|
provider_quota_writer: None,
|
||||||
usage_reader: None,
|
usage_reader: None,
|
||||||
@@ -333,6 +343,18 @@ impl GatewayDataState {
|
|||||||
self
|
self
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
pub(crate) fn with_pool_score_repository_for_tests<T>(mut self, repository: Arc<T>) -> Self
|
||||||
|
where
|
||||||
|
T: PoolMemberScoreRepository + 'static,
|
||||||
|
{
|
||||||
|
let pool_score_reader: Arc<dyn PoolScoreReadRepository> = repository.clone();
|
||||||
|
let pool_score_writer: Arc<dyn PoolMemberScoreWriteRepository> = repository;
|
||||||
|
self.pool_score_reader = Some(pool_score_reader);
|
||||||
|
self.pool_score_writer = Some(pool_score_writer);
|
||||||
|
self
|
||||||
|
}
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
pub(crate) fn with_provider_catalog_and_request_candidate_reader_for_tests(
|
pub(crate) fn with_provider_catalog_and_request_candidate_reader_for_tests(
|
||||||
provider_catalog_repository: Arc<dyn ProviderCatalogReadRepository>,
|
provider_catalog_repository: Arc<dyn ProviderCatalogReadRepository>,
|
||||||
@@ -363,6 +385,8 @@ impl GatewayDataState {
|
|||||||
request_candidate_writer: None,
|
request_candidate_writer: None,
|
||||||
provider_catalog_reader: Some(provider_catalog_repository),
|
provider_catalog_reader: Some(provider_catalog_repository),
|
||||||
provider_catalog_writer: None,
|
provider_catalog_writer: None,
|
||||||
|
pool_score_reader: None,
|
||||||
|
pool_score_writer: None,
|
||||||
provider_quota_reader: None,
|
provider_quota_reader: None,
|
||||||
provider_quota_writer: None,
|
provider_quota_writer: None,
|
||||||
usage_reader: None,
|
usage_reader: None,
|
||||||
@@ -419,6 +443,8 @@ impl GatewayDataState {
|
|||||||
request_candidate_writer: None,
|
request_candidate_writer: None,
|
||||||
provider_catalog_reader: Some(provider_catalog_repository),
|
provider_catalog_reader: Some(provider_catalog_repository),
|
||||||
provider_catalog_writer: None,
|
provider_catalog_writer: None,
|
||||||
|
pool_score_reader: None,
|
||||||
|
pool_score_writer: None,
|
||||||
provider_quota_reader: Some(provider_quota_reader),
|
provider_quota_reader: Some(provider_quota_reader),
|
||||||
provider_quota_writer: Some(provider_quota_writer),
|
provider_quota_writer: Some(provider_quota_writer),
|
||||||
usage_reader: None,
|
usage_reader: None,
|
||||||
@@ -484,6 +510,8 @@ impl GatewayDataState {
|
|||||||
request_candidate_writer: None,
|
request_candidate_writer: None,
|
||||||
provider_catalog_reader: Some(provider_catalog_reader),
|
provider_catalog_reader: Some(provider_catalog_reader),
|
||||||
provider_catalog_writer: Some(provider_catalog_writer),
|
provider_catalog_writer: Some(provider_catalog_writer),
|
||||||
|
pool_score_reader: None,
|
||||||
|
pool_score_writer: None,
|
||||||
provider_quota_reader: Some(provider_quota_reader),
|
provider_quota_reader: Some(provider_quota_reader),
|
||||||
provider_quota_writer: Some(provider_quota_writer),
|
provider_quota_writer: Some(provider_quota_writer),
|
||||||
usage_reader: None,
|
usage_reader: None,
|
||||||
@@ -531,6 +559,8 @@ impl GatewayDataState {
|
|||||||
request_candidate_writer: None,
|
request_candidate_writer: None,
|
||||||
provider_catalog_reader: None,
|
provider_catalog_reader: None,
|
||||||
provider_catalog_writer: None,
|
provider_catalog_writer: None,
|
||||||
|
pool_score_reader: None,
|
||||||
|
pool_score_writer: None,
|
||||||
provider_quota_reader: None,
|
provider_quota_reader: None,
|
||||||
provider_quota_writer: None,
|
provider_quota_writer: None,
|
||||||
usage_reader: None,
|
usage_reader: None,
|
||||||
@@ -579,6 +609,8 @@ impl GatewayDataState {
|
|||||||
request_candidate_writer: None,
|
request_candidate_writer: None,
|
||||||
provider_catalog_reader: Some(provider_catalog_repository),
|
provider_catalog_reader: Some(provider_catalog_repository),
|
||||||
provider_catalog_writer: None,
|
provider_catalog_writer: None,
|
||||||
|
pool_score_reader: None,
|
||||||
|
pool_score_writer: None,
|
||||||
provider_quota_reader: None,
|
provider_quota_reader: None,
|
||||||
provider_quota_writer: None,
|
provider_quota_writer: None,
|
||||||
usage_reader: None,
|
usage_reader: None,
|
||||||
@@ -638,6 +670,8 @@ impl GatewayDataState {
|
|||||||
request_candidate_writer: Some(request_candidate_writer),
|
request_candidate_writer: Some(request_candidate_writer),
|
||||||
provider_catalog_reader: None,
|
provider_catalog_reader: None,
|
||||||
provider_catalog_writer: None,
|
provider_catalog_writer: None,
|
||||||
|
pool_score_reader: None,
|
||||||
|
pool_score_writer: None,
|
||||||
provider_quota_reader: None,
|
provider_quota_reader: None,
|
||||||
provider_quota_writer: None,
|
provider_quota_writer: None,
|
||||||
usage_reader: Some(usage_reader),
|
usage_reader: Some(usage_reader),
|
||||||
@@ -699,6 +733,8 @@ impl GatewayDataState {
|
|||||||
request_candidate_writer: Some(request_candidate_writer),
|
request_candidate_writer: Some(request_candidate_writer),
|
||||||
provider_catalog_reader: None,
|
provider_catalog_reader: None,
|
||||||
provider_catalog_writer: None,
|
provider_catalog_writer: None,
|
||||||
|
pool_score_reader: None,
|
||||||
|
pool_score_writer: None,
|
||||||
provider_quota_reader: None,
|
provider_quota_reader: None,
|
||||||
provider_quota_writer: None,
|
provider_quota_writer: None,
|
||||||
usage_reader: None,
|
usage_reader: None,
|
||||||
@@ -744,6 +780,8 @@ impl GatewayDataState {
|
|||||||
request_candidate_writer: None,
|
request_candidate_writer: None,
|
||||||
provider_catalog_reader: None,
|
provider_catalog_reader: None,
|
||||||
provider_catalog_writer: None,
|
provider_catalog_writer: None,
|
||||||
|
pool_score_reader: None,
|
||||||
|
pool_score_writer: None,
|
||||||
provider_quota_reader: None,
|
provider_quota_reader: None,
|
||||||
provider_quota_writer: None,
|
provider_quota_writer: None,
|
||||||
usage_reader: Some(repository),
|
usage_reader: Some(repository),
|
||||||
@@ -804,6 +842,8 @@ impl GatewayDataState {
|
|||||||
request_candidate_writer: None,
|
request_candidate_writer: None,
|
||||||
provider_catalog_reader: None,
|
provider_catalog_reader: None,
|
||||||
provider_catalog_writer: None,
|
provider_catalog_writer: None,
|
||||||
|
pool_score_reader: None,
|
||||||
|
pool_score_writer: None,
|
||||||
provider_quota_reader: None,
|
provider_quota_reader: None,
|
||||||
provider_quota_writer: None,
|
provider_quota_writer: None,
|
||||||
usage_reader: None,
|
usage_reader: None,
|
||||||
@@ -857,6 +897,8 @@ impl GatewayDataState {
|
|||||||
request_candidate_writer: None,
|
request_candidate_writer: None,
|
||||||
provider_catalog_reader: None,
|
provider_catalog_reader: None,
|
||||||
provider_catalog_writer: None,
|
provider_catalog_writer: None,
|
||||||
|
pool_score_reader: None,
|
||||||
|
pool_score_writer: None,
|
||||||
provider_quota_reader: None,
|
provider_quota_reader: None,
|
||||||
provider_quota_writer: None,
|
provider_quota_writer: None,
|
||||||
usage_reader: None,
|
usage_reader: None,
|
||||||
@@ -915,6 +957,8 @@ impl GatewayDataState {
|
|||||||
request_candidate_writer: None,
|
request_candidate_writer: None,
|
||||||
provider_catalog_reader: None,
|
provider_catalog_reader: None,
|
||||||
provider_catalog_writer: None,
|
provider_catalog_writer: None,
|
||||||
|
pool_score_reader: None,
|
||||||
|
pool_score_writer: None,
|
||||||
provider_quota_reader: None,
|
provider_quota_reader: None,
|
||||||
provider_quota_writer: None,
|
provider_quota_writer: None,
|
||||||
usage_reader: Some(usage_reader),
|
usage_reader: Some(usage_reader),
|
||||||
@@ -974,6 +1018,8 @@ impl GatewayDataState {
|
|||||||
request_candidate_writer: None,
|
request_candidate_writer: None,
|
||||||
provider_catalog_reader: None,
|
provider_catalog_reader: None,
|
||||||
provider_catalog_writer: None,
|
provider_catalog_writer: None,
|
||||||
|
pool_score_reader: None,
|
||||||
|
pool_score_writer: None,
|
||||||
provider_quota_reader: None,
|
provider_quota_reader: None,
|
||||||
provider_quota_writer: None,
|
provider_quota_writer: None,
|
||||||
usage_reader: None,
|
usage_reader: None,
|
||||||
@@ -1032,6 +1078,8 @@ impl GatewayDataState {
|
|||||||
request_candidate_writer: None,
|
request_candidate_writer: None,
|
||||||
provider_catalog_reader: Some(provider_catalog_reader),
|
provider_catalog_reader: Some(provider_catalog_reader),
|
||||||
provider_catalog_writer: Some(provider_catalog_writer),
|
provider_catalog_writer: Some(provider_catalog_writer),
|
||||||
|
pool_score_reader: None,
|
||||||
|
pool_score_writer: None,
|
||||||
provider_quota_reader: None,
|
provider_quota_reader: None,
|
||||||
provider_quota_writer: None,
|
provider_quota_writer: None,
|
||||||
usage_reader: Some(usage_reader),
|
usage_reader: Some(usage_reader),
|
||||||
@@ -1079,6 +1127,8 @@ impl GatewayDataState {
|
|||||||
request_candidate_writer: None,
|
request_candidate_writer: None,
|
||||||
provider_catalog_reader: None,
|
provider_catalog_reader: None,
|
||||||
provider_catalog_writer: None,
|
provider_catalog_writer: None,
|
||||||
|
pool_score_reader: None,
|
||||||
|
pool_score_writer: None,
|
||||||
provider_quota_reader: None,
|
provider_quota_reader: None,
|
||||||
provider_quota_writer: None,
|
provider_quota_writer: None,
|
||||||
usage_reader: None,
|
usage_reader: None,
|
||||||
@@ -1126,6 +1176,8 @@ impl GatewayDataState {
|
|||||||
request_candidate_writer: None,
|
request_candidate_writer: None,
|
||||||
provider_catalog_reader: None,
|
provider_catalog_reader: None,
|
||||||
provider_catalog_writer: None,
|
provider_catalog_writer: None,
|
||||||
|
pool_score_reader: None,
|
||||||
|
pool_score_writer: None,
|
||||||
provider_quota_reader: None,
|
provider_quota_reader: None,
|
||||||
provider_quota_writer: None,
|
provider_quota_writer: None,
|
||||||
usage_reader: None,
|
usage_reader: None,
|
||||||
@@ -1185,6 +1237,8 @@ impl GatewayDataState {
|
|||||||
request_candidate_writer: None,
|
request_candidate_writer: None,
|
||||||
provider_catalog_reader: None,
|
provider_catalog_reader: None,
|
||||||
provider_catalog_writer: None,
|
provider_catalog_writer: None,
|
||||||
|
pool_score_reader: None,
|
||||||
|
pool_score_writer: None,
|
||||||
provider_quota_reader: None,
|
provider_quota_reader: None,
|
||||||
provider_quota_writer: None,
|
provider_quota_writer: None,
|
||||||
usage_reader: None,
|
usage_reader: None,
|
||||||
@@ -1249,6 +1303,8 @@ impl GatewayDataState {
|
|||||||
request_candidate_writer: None,
|
request_candidate_writer: None,
|
||||||
provider_catalog_reader: None,
|
provider_catalog_reader: None,
|
||||||
provider_catalog_writer: None,
|
provider_catalog_writer: None,
|
||||||
|
pool_score_reader: None,
|
||||||
|
pool_score_writer: None,
|
||||||
provider_quota_reader: None,
|
provider_quota_reader: None,
|
||||||
provider_quota_writer: None,
|
provider_quota_writer: None,
|
||||||
usage_reader: None,
|
usage_reader: None,
|
||||||
@@ -1296,6 +1352,8 @@ impl GatewayDataState {
|
|||||||
request_candidate_writer: None,
|
request_candidate_writer: None,
|
||||||
provider_catalog_reader: None,
|
provider_catalog_reader: None,
|
||||||
provider_catalog_writer: None,
|
provider_catalog_writer: None,
|
||||||
|
pool_score_reader: None,
|
||||||
|
pool_score_writer: None,
|
||||||
provider_quota_reader: None,
|
provider_quota_reader: None,
|
||||||
provider_quota_writer: None,
|
provider_quota_writer: None,
|
||||||
usage_reader: None,
|
usage_reader: None,
|
||||||
@@ -1348,6 +1406,8 @@ impl GatewayDataState {
|
|||||||
request_candidate_writer: None,
|
request_candidate_writer: None,
|
||||||
provider_catalog_reader: None,
|
provider_catalog_reader: None,
|
||||||
provider_catalog_writer: None,
|
provider_catalog_writer: None,
|
||||||
|
pool_score_reader: None,
|
||||||
|
pool_score_writer: None,
|
||||||
provider_quota_reader: None,
|
provider_quota_reader: None,
|
||||||
provider_quota_writer: None,
|
provider_quota_writer: None,
|
||||||
usage_reader: None,
|
usage_reader: None,
|
||||||
@@ -1417,6 +1477,8 @@ impl GatewayDataState {
|
|||||||
request_candidate_writer: None,
|
request_candidate_writer: None,
|
||||||
provider_catalog_reader: None,
|
provider_catalog_reader: None,
|
||||||
provider_catalog_writer: None,
|
provider_catalog_writer: None,
|
||||||
|
pool_score_reader: None,
|
||||||
|
pool_score_writer: None,
|
||||||
provider_quota_reader: None,
|
provider_quota_reader: None,
|
||||||
provider_quota_writer: None,
|
provider_quota_writer: None,
|
||||||
usage_reader: None,
|
usage_reader: None,
|
||||||
@@ -1481,6 +1543,8 @@ impl GatewayDataState {
|
|||||||
request_candidate_writer: None,
|
request_candidate_writer: None,
|
||||||
provider_catalog_reader: None,
|
provider_catalog_reader: None,
|
||||||
provider_catalog_writer: None,
|
provider_catalog_writer: None,
|
||||||
|
pool_score_reader: None,
|
||||||
|
pool_score_writer: None,
|
||||||
provider_quota_reader: None,
|
provider_quota_reader: None,
|
||||||
provider_quota_writer: None,
|
provider_quota_writer: None,
|
||||||
usage_reader: None,
|
usage_reader: None,
|
||||||
@@ -1529,6 +1593,8 @@ impl GatewayDataState {
|
|||||||
request_candidate_writer: None,
|
request_candidate_writer: None,
|
||||||
provider_catalog_reader: Some(provider_catalog_repository),
|
provider_catalog_reader: Some(provider_catalog_repository),
|
||||||
provider_catalog_writer: None,
|
provider_catalog_writer: None,
|
||||||
|
pool_score_reader: None,
|
||||||
|
pool_score_writer: None,
|
||||||
provider_quota_reader: None,
|
provider_quota_reader: None,
|
||||||
provider_quota_writer: None,
|
provider_quota_writer: None,
|
||||||
usage_reader: None,
|
usage_reader: None,
|
||||||
@@ -1577,6 +1643,8 @@ impl GatewayDataState {
|
|||||||
request_candidate_writer: None,
|
request_candidate_writer: None,
|
||||||
provider_catalog_reader: Some(repository),
|
provider_catalog_reader: Some(repository),
|
||||||
provider_catalog_writer: None,
|
provider_catalog_writer: None,
|
||||||
|
pool_score_reader: None,
|
||||||
|
pool_score_writer: None,
|
||||||
provider_quota_reader: None,
|
provider_quota_reader: None,
|
||||||
provider_quota_writer: None,
|
provider_quota_writer: None,
|
||||||
usage_reader: None,
|
usage_reader: None,
|
||||||
@@ -1627,6 +1695,8 @@ impl GatewayDataState {
|
|||||||
request_candidate_writer: None,
|
request_candidate_writer: None,
|
||||||
provider_catalog_reader: Some(provider_catalog_repository),
|
provider_catalog_reader: Some(provider_catalog_repository),
|
||||||
provider_catalog_writer: None,
|
provider_catalog_writer: None,
|
||||||
|
pool_score_reader: None,
|
||||||
|
pool_score_writer: None,
|
||||||
provider_quota_reader: None,
|
provider_quota_reader: None,
|
||||||
provider_quota_writer: None,
|
provider_quota_writer: None,
|
||||||
usage_reader: Some(usage_repository),
|
usage_reader: Some(usage_repository),
|
||||||
@@ -1675,6 +1745,8 @@ impl GatewayDataState {
|
|||||||
request_candidate_writer: None,
|
request_candidate_writer: None,
|
||||||
provider_catalog_reader: None,
|
provider_catalog_reader: None,
|
||||||
provider_catalog_writer: None,
|
provider_catalog_writer: None,
|
||||||
|
pool_score_reader: None,
|
||||||
|
pool_score_writer: None,
|
||||||
provider_quota_reader: None,
|
provider_quota_reader: None,
|
||||||
provider_quota_writer: None,
|
provider_quota_writer: None,
|
||||||
usage_reader: None,
|
usage_reader: None,
|
||||||
@@ -1723,6 +1795,8 @@ impl GatewayDataState {
|
|||||||
request_candidate_writer: None,
|
request_candidate_writer: None,
|
||||||
provider_catalog_reader: None,
|
provider_catalog_reader: None,
|
||||||
provider_catalog_writer: None,
|
provider_catalog_writer: None,
|
||||||
|
pool_score_reader: None,
|
||||||
|
pool_score_writer: None,
|
||||||
provider_quota_reader: None,
|
provider_quota_reader: None,
|
||||||
provider_quota_writer: None,
|
provider_quota_writer: None,
|
||||||
usage_reader: None,
|
usage_reader: None,
|
||||||
@@ -1779,6 +1853,8 @@ impl GatewayDataState {
|
|||||||
request_candidate_writer: None,
|
request_candidate_writer: None,
|
||||||
provider_catalog_reader: None,
|
provider_catalog_reader: None,
|
||||||
provider_catalog_writer: None,
|
provider_catalog_writer: None,
|
||||||
|
pool_score_reader: None,
|
||||||
|
pool_score_writer: None,
|
||||||
provider_quota_reader: Some(provider_quota_reader),
|
provider_quota_reader: Some(provider_quota_reader),
|
||||||
provider_quota_writer: Some(provider_quota_writer),
|
provider_quota_writer: Some(provider_quota_writer),
|
||||||
usage_reader: None,
|
usage_reader: None,
|
||||||
@@ -1836,6 +1912,8 @@ impl GatewayDataState {
|
|||||||
request_candidate_writer: None,
|
request_candidate_writer: None,
|
||||||
provider_catalog_reader: None,
|
provider_catalog_reader: None,
|
||||||
provider_catalog_writer: None,
|
provider_catalog_writer: None,
|
||||||
|
pool_score_reader: None,
|
||||||
|
pool_score_writer: None,
|
||||||
provider_quota_reader: Some(provider_quota_reader),
|
provider_quota_reader: Some(provider_quota_reader),
|
||||||
provider_quota_writer: Some(provider_quota_writer),
|
provider_quota_writer: Some(provider_quota_writer),
|
||||||
usage_reader: None,
|
usage_reader: None,
|
||||||
@@ -1896,6 +1974,8 @@ impl GatewayDataState {
|
|||||||
request_candidate_writer: None,
|
request_candidate_writer: None,
|
||||||
provider_catalog_reader: Some(provider_catalog_repository),
|
provider_catalog_reader: Some(provider_catalog_repository),
|
||||||
provider_catalog_writer: None,
|
provider_catalog_writer: None,
|
||||||
|
pool_score_reader: None,
|
||||||
|
pool_score_writer: None,
|
||||||
provider_quota_reader: Some(provider_quota_reader),
|
provider_quota_reader: Some(provider_quota_reader),
|
||||||
provider_quota_writer: Some(provider_quota_writer),
|
provider_quota_writer: Some(provider_quota_writer),
|
||||||
usage_reader: None,
|
usage_reader: None,
|
||||||
@@ -1962,6 +2042,8 @@ impl GatewayDataState {
|
|||||||
request_candidate_writer: Some(request_candidate_writer),
|
request_candidate_writer: Some(request_candidate_writer),
|
||||||
provider_catalog_reader: Some(provider_catalog_reader),
|
provider_catalog_reader: Some(provider_catalog_reader),
|
||||||
provider_catalog_writer: Some(provider_catalog_writer),
|
provider_catalog_writer: Some(provider_catalog_writer),
|
||||||
|
pool_score_reader: None,
|
||||||
|
pool_score_writer: None,
|
||||||
provider_quota_reader: None,
|
provider_quota_reader: None,
|
||||||
provider_quota_writer: None,
|
provider_quota_writer: None,
|
||||||
usage_reader: None,
|
usage_reader: None,
|
||||||
@@ -2029,6 +2111,8 @@ impl GatewayDataState {
|
|||||||
request_candidate_writer: Some(request_candidate_writer),
|
request_candidate_writer: Some(request_candidate_writer),
|
||||||
provider_catalog_reader: Some(provider_catalog_reader),
|
provider_catalog_reader: Some(provider_catalog_reader),
|
||||||
provider_catalog_writer: Some(provider_catalog_writer),
|
provider_catalog_writer: Some(provider_catalog_writer),
|
||||||
|
pool_score_reader: None,
|
||||||
|
pool_score_writer: None,
|
||||||
provider_quota_reader: None,
|
provider_quota_reader: None,
|
||||||
provider_quota_writer: None,
|
provider_quota_writer: None,
|
||||||
usage_reader: None,
|
usage_reader: None,
|
||||||
@@ -2100,6 +2184,8 @@ impl GatewayDataState {
|
|||||||
request_candidate_writer: Some(request_candidate_writer),
|
request_candidate_writer: Some(request_candidate_writer),
|
||||||
provider_catalog_reader: Some(provider_catalog_reader),
|
provider_catalog_reader: Some(provider_catalog_reader),
|
||||||
provider_catalog_writer: Some(provider_catalog_writer),
|
provider_catalog_writer: Some(provider_catalog_writer),
|
||||||
|
pool_score_reader: None,
|
||||||
|
pool_score_writer: None,
|
||||||
provider_quota_reader: None,
|
provider_quota_reader: None,
|
||||||
provider_quota_writer: None,
|
provider_quota_writer: None,
|
||||||
usage_reader: Some(usage_reader),
|
usage_reader: Some(usage_reader),
|
||||||
@@ -2178,6 +2264,8 @@ impl GatewayDataState {
|
|||||||
request_candidate_writer: Some(request_candidate_writer),
|
request_candidate_writer: Some(request_candidate_writer),
|
||||||
provider_catalog_reader: Some(provider_catalog_reader),
|
provider_catalog_reader: Some(provider_catalog_reader),
|
||||||
provider_catalog_writer: Some(provider_catalog_writer),
|
provider_catalog_writer: Some(provider_catalog_writer),
|
||||||
|
pool_score_reader: None,
|
||||||
|
pool_score_writer: None,
|
||||||
provider_quota_reader: None,
|
provider_quota_reader: None,
|
||||||
provider_quota_writer: None,
|
provider_quota_writer: None,
|
||||||
usage_reader: Some(usage_reader),
|
usage_reader: Some(usage_reader),
|
||||||
@@ -2238,6 +2326,8 @@ impl GatewayDataState {
|
|||||||
request_candidate_writer: None,
|
request_candidate_writer: None,
|
||||||
provider_catalog_reader: Some(provider_catalog_repository),
|
provider_catalog_reader: Some(provider_catalog_repository),
|
||||||
provider_catalog_writer: None,
|
provider_catalog_writer: None,
|
||||||
|
pool_score_reader: None,
|
||||||
|
pool_score_writer: None,
|
||||||
provider_quota_reader: Some(provider_quota_reader),
|
provider_quota_reader: Some(provider_quota_reader),
|
||||||
provider_quota_writer: Some(provider_quota_writer),
|
provider_quota_writer: Some(provider_quota_writer),
|
||||||
usage_reader: None,
|
usage_reader: None,
|
||||||
@@ -2289,6 +2379,8 @@ impl GatewayDataState {
|
|||||||
request_candidate_writer: None,
|
request_candidate_writer: None,
|
||||||
provider_catalog_reader: None,
|
provider_catalog_reader: None,
|
||||||
provider_catalog_writer: None,
|
provider_catalog_writer: None,
|
||||||
|
pool_score_reader: None,
|
||||||
|
pool_score_writer: None,
|
||||||
provider_quota_reader: None,
|
provider_quota_reader: None,
|
||||||
provider_quota_writer: None,
|
provider_quota_writer: None,
|
||||||
usage_reader: Some(usage_reader),
|
usage_reader: Some(usage_reader),
|
||||||
@@ -2336,6 +2428,8 @@ impl GatewayDataState {
|
|||||||
request_candidate_writer: None,
|
request_candidate_writer: None,
|
||||||
provider_catalog_reader: None,
|
provider_catalog_reader: None,
|
||||||
provider_catalog_writer: None,
|
provider_catalog_writer: None,
|
||||||
|
pool_score_reader: None,
|
||||||
|
pool_score_writer: None,
|
||||||
provider_quota_reader: None,
|
provider_quota_reader: None,
|
||||||
provider_quota_writer: None,
|
provider_quota_writer: None,
|
||||||
usage_reader: None,
|
usage_reader: None,
|
||||||
@@ -2389,6 +2483,8 @@ impl GatewayDataState {
|
|||||||
request_candidate_writer: None,
|
request_candidate_writer: None,
|
||||||
provider_catalog_reader: None,
|
provider_catalog_reader: None,
|
||||||
provider_catalog_writer: None,
|
provider_catalog_writer: None,
|
||||||
|
pool_score_reader: None,
|
||||||
|
pool_score_writer: None,
|
||||||
provider_quota_reader: None,
|
provider_quota_reader: None,
|
||||||
provider_quota_writer: None,
|
provider_quota_writer: None,
|
||||||
usage_reader: None,
|
usage_reader: None,
|
||||||
@@ -2446,6 +2542,8 @@ impl GatewayDataState {
|
|||||||
request_candidate_writer: None,
|
request_candidate_writer: None,
|
||||||
provider_catalog_reader: None,
|
provider_catalog_reader: None,
|
||||||
provider_catalog_writer: None,
|
provider_catalog_writer: None,
|
||||||
|
pool_score_reader: None,
|
||||||
|
pool_score_writer: None,
|
||||||
provider_quota_reader: None,
|
provider_quota_reader: None,
|
||||||
provider_quota_writer: None,
|
provider_quota_writer: None,
|
||||||
usage_reader: Some(usage_reader),
|
usage_reader: Some(usage_reader),
|
||||||
@@ -2504,6 +2602,8 @@ impl GatewayDataState {
|
|||||||
request_candidate_writer: None,
|
request_candidate_writer: None,
|
||||||
provider_catalog_reader: None,
|
provider_catalog_reader: None,
|
||||||
provider_catalog_writer: None,
|
provider_catalog_writer: None,
|
||||||
|
pool_score_reader: None,
|
||||||
|
pool_score_writer: None,
|
||||||
provider_quota_reader: None,
|
provider_quota_reader: None,
|
||||||
provider_quota_writer: None,
|
provider_quota_writer: None,
|
||||||
usage_reader: Some(usage_reader),
|
usage_reader: Some(usage_reader),
|
||||||
@@ -2555,6 +2655,8 @@ impl GatewayDataState {
|
|||||||
request_candidate_writer: None,
|
request_candidate_writer: None,
|
||||||
provider_catalog_reader: None,
|
provider_catalog_reader: None,
|
||||||
provider_catalog_writer: None,
|
provider_catalog_writer: None,
|
||||||
|
pool_score_reader: None,
|
||||||
|
pool_score_writer: None,
|
||||||
provider_quota_reader: Some(provider_quota_reader),
|
provider_quota_reader: Some(provider_quota_reader),
|
||||||
provider_quota_writer: Some(provider_quota_writer),
|
provider_quota_writer: Some(provider_quota_writer),
|
||||||
usage_reader: None,
|
usage_reader: None,
|
||||||
|
|||||||
@@ -38,6 +38,8 @@ impl GatewayDataState {
|
|||||||
request_candidate_writer: None,
|
request_candidate_writer: None,
|
||||||
provider_catalog_reader: None,
|
provider_catalog_reader: None,
|
||||||
provider_catalog_writer: None,
|
provider_catalog_writer: None,
|
||||||
|
pool_score_reader: None,
|
||||||
|
pool_score_writer: None,
|
||||||
provider_quota_reader: None,
|
provider_quota_reader: None,
|
||||||
provider_quota_writer: None,
|
provider_quota_writer: None,
|
||||||
usage_reader: None,
|
usage_reader: None,
|
||||||
@@ -92,6 +94,8 @@ impl GatewayDataState {
|
|||||||
request_candidate_writer: None,
|
request_candidate_writer: None,
|
||||||
provider_catalog_reader: None,
|
provider_catalog_reader: None,
|
||||||
provider_catalog_writer: None,
|
provider_catalog_writer: None,
|
||||||
|
pool_score_reader: None,
|
||||||
|
pool_score_writer: None,
|
||||||
provider_quota_reader: None,
|
provider_quota_reader: None,
|
||||||
provider_quota_writer: None,
|
provider_quota_writer: None,
|
||||||
usage_reader: None,
|
usage_reader: None,
|
||||||
@@ -143,6 +147,8 @@ impl GatewayDataState {
|
|||||||
request_candidate_writer: None,
|
request_candidate_writer: None,
|
||||||
provider_catalog_reader: None,
|
provider_catalog_reader: None,
|
||||||
provider_catalog_writer: None,
|
provider_catalog_writer: None,
|
||||||
|
pool_score_reader: None,
|
||||||
|
pool_score_writer: None,
|
||||||
provider_quota_reader: None,
|
provider_quota_reader: None,
|
||||||
provider_quota_writer: None,
|
provider_quota_writer: None,
|
||||||
usage_reader: None,
|
usage_reader: None,
|
||||||
@@ -198,6 +204,8 @@ impl GatewayDataState {
|
|||||||
request_candidate_writer: None,
|
request_candidate_writer: None,
|
||||||
provider_catalog_reader: Some(provider_catalog_repository),
|
provider_catalog_reader: Some(provider_catalog_repository),
|
||||||
provider_catalog_writer: None,
|
provider_catalog_writer: None,
|
||||||
|
pool_score_reader: None,
|
||||||
|
pool_score_writer: None,
|
||||||
provider_quota_reader: None,
|
provider_quota_reader: None,
|
||||||
provider_quota_writer: None,
|
provider_quota_writer: None,
|
||||||
usage_reader: None,
|
usage_reader: None,
|
||||||
@@ -257,6 +265,8 @@ impl GatewayDataState {
|
|||||||
request_candidate_writer: Some(request_candidate_writer),
|
request_candidate_writer: Some(request_candidate_writer),
|
||||||
provider_catalog_reader: None,
|
provider_catalog_reader: None,
|
||||||
provider_catalog_writer: None,
|
provider_catalog_writer: None,
|
||||||
|
pool_score_reader: None,
|
||||||
|
pool_score_writer: None,
|
||||||
provider_quota_reader: None,
|
provider_quota_reader: None,
|
||||||
provider_quota_writer: None,
|
provider_quota_writer: None,
|
||||||
usage_reader: None,
|
usage_reader: None,
|
||||||
@@ -325,6 +335,8 @@ impl GatewayDataState {
|
|||||||
request_candidate_writer: Some(request_candidate_writer),
|
request_candidate_writer: Some(request_candidate_writer),
|
||||||
provider_catalog_reader: Some(provider_catalog_reader),
|
provider_catalog_reader: Some(provider_catalog_reader),
|
||||||
provider_catalog_writer: None,
|
provider_catalog_writer: None,
|
||||||
|
pool_score_reader: None,
|
||||||
|
pool_score_writer: None,
|
||||||
provider_quota_reader: None,
|
provider_quota_reader: None,
|
||||||
provider_quota_writer: None,
|
provider_quota_writer: None,
|
||||||
usage_reader: None,
|
usage_reader: None,
|
||||||
|
|||||||
@@ -1,6 +1,8 @@
|
|||||||
|
use crate::handlers::admin::admin_provider_pool_config;
|
||||||
use crate::handlers::admin::provider::shared::paths::admin_provider_id_for_keys;
|
use crate::handlers::admin::provider::shared::paths::admin_provider_id_for_keys;
|
||||||
use crate::handlers::admin::provider::shared::payloads::AdminProviderKeyCreateRequest;
|
use crate::handlers::admin::provider::shared::payloads::AdminProviderKeyCreateRequest;
|
||||||
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
|
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
|
||||||
|
use crate::maintenance::ensure_provider_key_pool_scores_for_keys;
|
||||||
use crate::provider_key_auth::provider_key_effective_api_formats;
|
use crate::provider_key_auth::provider_key_effective_api_formats;
|
||||||
use crate::{model_fetch::perform_model_fetch_for_key, GatewayError};
|
use crate::{model_fetch::perform_model_fetch_for_key, GatewayError};
|
||||||
use axum::{
|
use axum::{
|
||||||
@@ -98,6 +100,27 @@ pub(super) async fn maybe_handle(
|
|||||||
let endpoints = state
|
let endpoints = state
|
||||||
.list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider.id))
|
.list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider.id))
|
||||||
.await?;
|
.await?;
|
||||||
|
if let Some(pool_config) = admin_provider_pool_config(&provider) {
|
||||||
|
let score_ensure_budget = (pool_config.score_fallback_scan_limit as usize).clamp(1, 50_000);
|
||||||
|
if let Err(err) = ensure_provider_key_pool_scores_for_keys(
|
||||||
|
state.as_ref(),
|
||||||
|
&provider,
|
||||||
|
&pool_config,
|
||||||
|
&endpoints,
|
||||||
|
std::slice::from_ref(&created),
|
||||||
|
now_unix_secs,
|
||||||
|
score_ensure_budget,
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
tracing::debug!(
|
||||||
|
provider_id = %provider.id,
|
||||||
|
key_id = %created.id,
|
||||||
|
error = ?err,
|
||||||
|
"gateway admin provider key create: failed to seed pool score rows"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
let api_formats =
|
let api_formats =
|
||||||
provider_key_effective_api_formats(&created, &provider.provider_type, &endpoints);
|
provider_key_effective_api_formats(&created, &provider.provider_type, &endpoints);
|
||||||
|
|
||||||
|
|||||||
@@ -1,6 +1,8 @@
|
|||||||
|
use crate::handlers::admin::admin_provider_pool_config;
|
||||||
use crate::handlers::admin::provider::shared::paths::admin_update_key_id;
|
use crate::handlers::admin::provider::shared::paths::admin_update_key_id;
|
||||||
use crate::handlers::admin::provider::shared::payloads::AdminProviderKeyUpdatePatch;
|
use crate::handlers::admin::provider::shared::payloads::AdminProviderKeyUpdatePatch;
|
||||||
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
|
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
|
||||||
|
use crate::maintenance::ensure_provider_key_pool_scores_for_keys;
|
||||||
use crate::provider_key_auth::provider_key_effective_api_formats;
|
use crate::provider_key_auth::provider_key_effective_api_formats;
|
||||||
use crate::{model_fetch::perform_model_fetch_for_key, GatewayError};
|
use crate::{model_fetch::perform_model_fetch_for_key, GatewayError};
|
||||||
use axum::{
|
use axum::{
|
||||||
@@ -121,6 +123,27 @@ pub(super) async fn maybe_handle(
|
|||||||
let endpoints = state
|
let endpoints = state
|
||||||
.list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider.id))
|
.list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider.id))
|
||||||
.await?;
|
.await?;
|
||||||
|
if let Some(pool_config) = admin_provider_pool_config(&provider) {
|
||||||
|
let score_ensure_budget = (pool_config.score_fallback_scan_limit as usize).clamp(1, 50_000);
|
||||||
|
if let Err(err) = ensure_provider_key_pool_scores_for_keys(
|
||||||
|
state.as_ref(),
|
||||||
|
&provider,
|
||||||
|
&pool_config,
|
||||||
|
&endpoints,
|
||||||
|
std::slice::from_ref(&updated),
|
||||||
|
now_unix_secs,
|
||||||
|
score_ensure_budget,
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
tracing::debug!(
|
||||||
|
provider_id = %provider.id,
|
||||||
|
key_id = %updated.id,
|
||||||
|
error = ?err,
|
||||||
|
"gateway admin provider key update: failed to seed pool score rows"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
let api_formats =
|
let api_formats =
|
||||||
provider_key_effective_api_formats(&updated, &provider.provider_type, &endpoints);
|
provider_key_effective_api_formats(&updated, &provider.provider_type, &endpoints);
|
||||||
|
|
||||||
|
|||||||
@@ -1,6 +1,7 @@
|
|||||||
use crate::handlers::admin::provider::shared::support::{
|
use crate::handlers::admin::provider::shared::support::{
|
||||||
AdminProviderPoolConfig, AdminProviderPoolSchedulingPreset, AdminProviderPoolUnschedulableRule,
|
AdminProviderPoolConfig, AdminProviderPoolSchedulingPreset, AdminProviderPoolUnschedulableRule,
|
||||||
};
|
};
|
||||||
|
use aether_ai_serving::{PoolMemberScoreRules, PoolMemberScoreWeights};
|
||||||
use serde_json::{Map, Value};
|
use serde_json::{Map, Value};
|
||||||
|
|
||||||
const POOL_ALLOWED_SCHEDULING_PRESETS: &[&str] = &[
|
const POOL_ALLOWED_SCHEDULING_PRESETS: &[&str] = &[
|
||||||
@@ -26,6 +27,123 @@ fn json_u64(value: &Value) -> Option<u64> {
|
|||||||
.or_else(|| value.as_i64().and_then(|raw| u64::try_from(raw).ok()))
|
.or_else(|| value.as_i64().and_then(|raw| u64::try_from(raw).ok()))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn json_f64(value: &Value) -> Option<f64> {
|
||||||
|
value.as_f64().or_else(|| {
|
||||||
|
value
|
||||||
|
.as_str()
|
||||||
|
.map(str::trim)
|
||||||
|
.filter(|value| !value.is_empty())
|
||||||
|
.and_then(|value| value.parse::<f64>().ok())
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn pool_score_weight(object: &Map<String, Value>, names: &[&str], current: f64) -> f64 {
|
||||||
|
names
|
||||||
|
.iter()
|
||||||
|
.find_map(|name| {
|
||||||
|
object
|
||||||
|
.get(*name)
|
||||||
|
.and_then(json_f64)
|
||||||
|
.filter(|value| value.is_finite() && *value >= 0.0)
|
||||||
|
})
|
||||||
|
.unwrap_or(current)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn parse_pool_score_weights(
|
||||||
|
raw_weights: Option<&Map<String, Value>>,
|
||||||
|
current: PoolMemberScoreWeights,
|
||||||
|
) -> PoolMemberScoreWeights {
|
||||||
|
let Some(raw_weights) = raw_weights else {
|
||||||
|
return current;
|
||||||
|
};
|
||||||
|
PoolMemberScoreWeights {
|
||||||
|
manual_priority: pool_score_weight(
|
||||||
|
raw_weights,
|
||||||
|
&["manual_priority", "priority", "internal_priority"],
|
||||||
|
current.manual_priority,
|
||||||
|
),
|
||||||
|
health: pool_score_weight(raw_weights, &["health"], current.health),
|
||||||
|
probe_freshness: pool_score_weight(
|
||||||
|
raw_weights,
|
||||||
|
&["probe_freshness", "freshness", "probe"],
|
||||||
|
current.probe_freshness,
|
||||||
|
),
|
||||||
|
quota_remaining: pool_score_weight(
|
||||||
|
raw_weights,
|
||||||
|
&["quota_remaining", "quota", "quota_available"],
|
||||||
|
current.quota_remaining,
|
||||||
|
),
|
||||||
|
latency: pool_score_weight(raw_weights, &["latency"], current.latency),
|
||||||
|
cost_lru: pool_score_weight(
|
||||||
|
raw_weights,
|
||||||
|
&["cost_lru", "cost_remaining", "cost", "lru"],
|
||||||
|
current.cost_lru,
|
||||||
|
),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn parse_pool_score_rules(pool_advanced: &Map<String, Value>) -> PoolMemberScoreRules {
|
||||||
|
let mut rules = PoolMemberScoreRules::default();
|
||||||
|
|
||||||
|
for key in ["score_weights", "pool_score_weights", "scoring_weights"] {
|
||||||
|
rules.weights = parse_pool_score_weights(
|
||||||
|
pool_advanced.get(key).and_then(Value::as_object),
|
||||||
|
rules.weights,
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
if let Some(score_rules) = pool_advanced
|
||||||
|
.get("score_rules")
|
||||||
|
.or_else(|| pool_advanced.get("pool_score_rules"))
|
||||||
|
.and_then(Value::as_object)
|
||||||
|
{
|
||||||
|
rules.weights = parse_pool_score_weights(
|
||||||
|
score_rules.get("weights").and_then(Value::as_object),
|
||||||
|
rules.weights,
|
||||||
|
);
|
||||||
|
if let Some(ttl_seconds) = score_rules
|
||||||
|
.get("probe_freshness_ttl_seconds")
|
||||||
|
.or_else(|| score_rules.get("score_probe_freshness_ttl_seconds"))
|
||||||
|
.and_then(json_u64)
|
||||||
|
.filter(|value| *value > 0)
|
||||||
|
{
|
||||||
|
rules.probe_freshness_ttl_seconds = ttl_seconds.min(7 * 24 * 3600);
|
||||||
|
}
|
||||||
|
if let Some(cap) = score_rules
|
||||||
|
.get("unschedulable_score_cap")
|
||||||
|
.or_else(|| score_rules.get("hard_state_score_cap"))
|
||||||
|
.and_then(json_f64)
|
||||||
|
.filter(|value| value.is_finite())
|
||||||
|
{
|
||||||
|
rules.unschedulable_score_cap = cap.clamp(0.0, 1.0);
|
||||||
|
}
|
||||||
|
if let Some(penalty) = score_rules
|
||||||
|
.get("probe_failure_penalty")
|
||||||
|
.and_then(json_f64)
|
||||||
|
.filter(|value| value.is_finite())
|
||||||
|
{
|
||||||
|
rules.probe_failure_penalty = penalty.clamp(0.0, 1.0);
|
||||||
|
}
|
||||||
|
if let Some(penalty) = score_rules
|
||||||
|
.get("request_failure_penalty")
|
||||||
|
.or_else(|| score_rules.get("runtime_failure_penalty"))
|
||||||
|
.and_then(json_f64)
|
||||||
|
.filter(|value| value.is_finite())
|
||||||
|
{
|
||||||
|
rules.request_failure_penalty = penalty.clamp(0.0, 1.0);
|
||||||
|
}
|
||||||
|
if let Some(threshold) = score_rules
|
||||||
|
.get("probe_failure_cooldown_threshold")
|
||||||
|
.or_else(|| score_rules.get("probe_failure_hard_state_threshold"))
|
||||||
|
.and_then(json_u64)
|
||||||
|
{
|
||||||
|
rules.probe_failure_cooldown_threshold = threshold.min(100);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
rules.effective()
|
||||||
|
}
|
||||||
|
|
||||||
fn normalize_pool_preset_mode(preset: &str, raw_mode: Option<&Value>) -> Option<String> {
|
fn normalize_pool_preset_mode(preset: &str, raw_mode: Option<&Value>) -> Option<String> {
|
||||||
match preset {
|
match preset {
|
||||||
"free_first" | "team_first" | "plus_first" | "pro_first" => {
|
"free_first" | "team_first" | "plus_first" | "pro_first" => {
|
||||||
@@ -270,6 +388,10 @@ pub(crate) fn admin_provider_pool_config_from_config_value(
|
|||||||
health_policy_enabled: true,
|
health_policy_enabled: true,
|
||||||
probing_enabled: false,
|
probing_enabled: false,
|
||||||
probing_interval_minutes: 10,
|
probing_interval_minutes: 10,
|
||||||
|
probe_concurrency: 4,
|
||||||
|
score_top_n: 128,
|
||||||
|
score_fallback_scan_limit: 1024,
|
||||||
|
score_rules: PoolMemberScoreRules::default(),
|
||||||
stream_timeout_threshold: 3,
|
stream_timeout_threshold: 3,
|
||||||
stream_timeout_window_seconds: 1800,
|
stream_timeout_window_seconds: 1800,
|
||||||
stream_timeout_cooldown_seconds: 300,
|
stream_timeout_cooldown_seconds: 300,
|
||||||
@@ -278,6 +400,7 @@ pub(crate) fn admin_provider_pool_config_from_config_value(
|
|||||||
|
|
||||||
let scheduling_presets = parse_pool_scheduling_presets(pool_advanced);
|
let scheduling_presets = parse_pool_scheduling_presets(pool_advanced);
|
||||||
let unschedulable_rules = parse_pool_unschedulable_rules(pool_advanced);
|
let unschedulable_rules = parse_pool_unschedulable_rules(pool_advanced);
|
||||||
|
let score_rules = parse_pool_score_rules(pool_advanced);
|
||||||
|
|
||||||
Some(AdminProviderPoolConfig {
|
Some(AdminProviderPoolConfig {
|
||||||
lru_enabled: admin_provider_pool_lru_enabled(&scheduling_presets),
|
lru_enabled: admin_provider_pool_lru_enabled(&scheduling_presets),
|
||||||
@@ -333,6 +456,25 @@ pub(crate) fn admin_provider_pool_config_from_config_value(
|
|||||||
.filter(|value| *value > 0)
|
.filter(|value| *value > 0)
|
||||||
.map(|value| value.min(1440))
|
.map(|value| value.min(1440))
|
||||||
.unwrap_or(10),
|
.unwrap_or(10),
|
||||||
|
probe_concurrency: pool_advanced
|
||||||
|
.get("probe_concurrency")
|
||||||
|
.and_then(json_u64)
|
||||||
|
.filter(|value| *value > 0)
|
||||||
|
.map(|value| value.min(64))
|
||||||
|
.unwrap_or(4),
|
||||||
|
score_top_n: pool_advanced
|
||||||
|
.get("score_top_n")
|
||||||
|
.and_then(json_u64)
|
||||||
|
.filter(|value| *value > 0)
|
||||||
|
.map(|value| value.min(4096))
|
||||||
|
.unwrap_or(128),
|
||||||
|
score_fallback_scan_limit: pool_advanced
|
||||||
|
.get("score_fallback_scan_limit")
|
||||||
|
.and_then(json_u64)
|
||||||
|
.filter(|value| *value > 0)
|
||||||
|
.map(|value| value.min(50_000))
|
||||||
|
.unwrap_or(1024),
|
||||||
|
score_rules,
|
||||||
stream_timeout_threshold: pool_advanced
|
stream_timeout_threshold: pool_advanced
|
||||||
.get("stream_timeout_threshold")
|
.get("stream_timeout_threshold")
|
||||||
.and_then(json_u64)
|
.and_then(json_u64)
|
||||||
@@ -405,6 +547,24 @@ mod tests {
|
|||||||
"health_policy_enabled": false,
|
"health_policy_enabled": false,
|
||||||
"probing_enabled": true,
|
"probing_enabled": true,
|
||||||
"probing_interval_minutes": 20,
|
"probing_interval_minutes": 20,
|
||||||
|
"probe_concurrency": 6,
|
||||||
|
"score_top_n": 256,
|
||||||
|
"score_fallback_scan_limit": 2048,
|
||||||
|
"score_rules": {
|
||||||
|
"weights": {
|
||||||
|
"manual_priority": 0.4,
|
||||||
|
"health": 0.2,
|
||||||
|
"probe_freshness": 0.2,
|
||||||
|
"quota_remaining": 0.1,
|
||||||
|
"latency": 0.05,
|
||||||
|
"cost_lru": 0.05
|
||||||
|
},
|
||||||
|
"probe_freshness_ttl_seconds": 1200,
|
||||||
|
"unschedulable_score_cap": 0.03,
|
||||||
|
"probe_failure_penalty": 0.08,
|
||||||
|
"request_failure_penalty": 0.01,
|
||||||
|
"probe_failure_cooldown_threshold": 2
|
||||||
|
},
|
||||||
"stream_timeout_threshold": 4,
|
"stream_timeout_threshold": 4,
|
||||||
"stream_timeout_window_seconds": 900,
|
"stream_timeout_window_seconds": 900,
|
||||||
"stream_timeout_cooldown_seconds": 180
|
"stream_timeout_cooldown_seconds": 180
|
||||||
@@ -424,6 +584,16 @@ mod tests {
|
|||||||
assert!(!config.health_policy_enabled);
|
assert!(!config.health_policy_enabled);
|
||||||
assert!(config.probing_enabled);
|
assert!(config.probing_enabled);
|
||||||
assert_eq!(config.probing_interval_minutes, 20);
|
assert_eq!(config.probing_interval_minutes, 20);
|
||||||
|
assert_eq!(config.probe_concurrency, 6);
|
||||||
|
assert_eq!(config.score_top_n, 256);
|
||||||
|
assert_eq!(config.score_fallback_scan_limit, 2048);
|
||||||
|
assert_eq!(config.score_rules.weights.manual_priority, 0.4);
|
||||||
|
assert_eq!(config.score_rules.weights.health, 0.2);
|
||||||
|
assert_eq!(config.score_rules.probe_freshness_ttl_seconds, 1200);
|
||||||
|
assert_eq!(config.score_rules.unschedulable_score_cap, 0.03);
|
||||||
|
assert_eq!(config.score_rules.probe_failure_penalty, 0.08);
|
||||||
|
assert_eq!(config.score_rules.request_failure_penalty, 0.01);
|
||||||
|
assert_eq!(config.score_rules.probe_failure_cooldown_threshold, 2);
|
||||||
assert_eq!(config.stream_timeout_threshold, 4);
|
assert_eq!(config.stream_timeout_threshold, 4);
|
||||||
assert_eq!(config.stream_timeout_window_seconds, 900);
|
assert_eq!(config.stream_timeout_window_seconds, 900);
|
||||||
assert_eq!(config.stream_timeout_cooldown_seconds, 180);
|
assert_eq!(config.stream_timeout_cooldown_seconds, 180);
|
||||||
@@ -462,6 +632,27 @@ mod tests {
|
|||||||
assert_eq!(config.sticky_session_ttl_seconds, 0);
|
assert_eq!(config.sticky_session_ttl_seconds, 0);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn parses_legacy_pool_score_weights_from_pool_advanced() {
|
||||||
|
let config = admin_provider_pool_config_from_config_value(Some(&json!({
|
||||||
|
"pool_advanced": {
|
||||||
|
"scoring_weights": {
|
||||||
|
"manual_priority": 0,
|
||||||
|
"health": 2,
|
||||||
|
"probe": 1,
|
||||||
|
"quota_remaining": 0,
|
||||||
|
"latency": 0,
|
||||||
|
"cost_remaining": 1
|
||||||
|
}
|
||||||
|
}
|
||||||
|
})))
|
||||||
|
.expect("pool config should parse");
|
||||||
|
|
||||||
|
assert_eq!(config.score_rules.weights.health, 0.5);
|
||||||
|
assert_eq!(config.score_rules.weights.probe_freshness, 0.25);
|
||||||
|
assert_eq!(config.score_rules.weights.cost_lru, 0.25);
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn parses_pool_config_from_generic_config_value() {
|
fn parses_pool_config_from_generic_config_value() {
|
||||||
let config = admin_provider_pool_config_from_config_value(Some(&json!({
|
let config = admin_provider_pool_config_from_config_value(Some(&json!({
|
||||||
|
|||||||
@@ -615,6 +615,10 @@ mod tests {
|
|||||||
health_policy_enabled: true,
|
health_policy_enabled: true,
|
||||||
probing_enabled: false,
|
probing_enabled: false,
|
||||||
probing_interval_minutes: 10,
|
probing_interval_minutes: 10,
|
||||||
|
probe_concurrency: 4,
|
||||||
|
score_top_n: 128,
|
||||||
|
score_fallback_scan_limit: 1024,
|
||||||
|
score_rules: aether_ai_serving::PoolMemberScoreRules::default(),
|
||||||
stream_timeout_threshold: 3,
|
stream_timeout_threshold: 3,
|
||||||
stream_timeout_window_seconds: 1800,
|
stream_timeout_window_seconds: 1800,
|
||||||
stream_timeout_cooldown_seconds: 300,
|
stream_timeout_cooldown_seconds: 300,
|
||||||
|
|||||||
@@ -24,6 +24,8 @@ mod read_overview;
|
|||||||
mod read_presets;
|
mod read_presets;
|
||||||
#[path = "read_routes/resolve_selection.rs"]
|
#[path = "read_routes/resolve_selection.rs"]
|
||||||
mod read_resolve_selection;
|
mod read_resolve_selection;
|
||||||
|
#[path = "read_routes/scores.rs"]
|
||||||
|
mod read_scores;
|
||||||
pub(crate) mod selection;
|
pub(crate) mod selection;
|
||||||
mod support;
|
mod support;
|
||||||
|
|
||||||
@@ -34,11 +36,11 @@ pub(crate) use self::batch_shared::{
|
|||||||
AdminPoolBatchImportRequest,
|
AdminPoolBatchImportRequest,
|
||||||
};
|
};
|
||||||
pub(crate) use self::support::{
|
pub(crate) use self::support::{
|
||||||
admin_pool_provider_id_from_path, parse_admin_pool_key_sort, parse_admin_pool_page,
|
admin_pool_provider_id_from_path, admin_pool_provider_id_from_scores_path,
|
||||||
parse_admin_pool_page_size, parse_admin_pool_quick_selectors, parse_admin_pool_search,
|
parse_admin_pool_key_sort, parse_admin_pool_page, parse_admin_pool_page_size,
|
||||||
parse_admin_pool_status_filter, AdminPoolKeySort, AdminPoolKeySortDirection,
|
parse_admin_pool_quick_selectors, parse_admin_pool_search, parse_admin_pool_status_filter,
|
||||||
AdminPoolKeySortField, AdminPoolResolveSelectionRequest,
|
AdminPoolKeySort, AdminPoolKeySortDirection, AdminPoolKeySortField,
|
||||||
ADMIN_POOL_BANNED_KEY_CLEANUP_EMPTY_MESSAGE,
|
AdminPoolResolveSelectionRequest, ADMIN_POOL_BANNED_KEY_CLEANUP_EMPTY_MESSAGE,
|
||||||
ADMIN_POOL_PROVIDER_CATALOG_READER_UNAVAILABLE_DETAIL,
|
ADMIN_POOL_PROVIDER_CATALOG_READER_UNAVAILABLE_DETAIL,
|
||||||
ADMIN_POOL_PROVIDER_CATALOG_WRITER_UNAVAILABLE_DETAIL,
|
ADMIN_POOL_PROVIDER_CATALOG_WRITER_UNAVAILABLE_DETAIL,
|
||||||
};
|
};
|
||||||
@@ -103,6 +105,11 @@ pub(crate) async fn maybe_build_local_admin_pool_response(
|
|||||||
read_keys::build_admin_pool_list_keys_response(state, request_context).await?,
|
read_keys::build_admin_pool_list_keys_response(state, request_context).await?,
|
||||||
));
|
));
|
||||||
}
|
}
|
||||||
|
Some("scores") => {
|
||||||
|
return Ok(Some(
|
||||||
|
read_scores::build_admin_pool_scores_response(state, request_context).await?,
|
||||||
|
));
|
||||||
|
}
|
||||||
Some("resolve_selection") => {
|
Some("resolve_selection") => {
|
||||||
return Ok(Some(
|
return Ok(Some(
|
||||||
read_resolve_selection::build_admin_pool_resolve_selection_response(
|
read_resolve_selection::build_admin_pool_resolve_selection_response(
|
||||||
|
|||||||
@@ -0,0 +1,196 @@
|
|||||||
|
use super::{
|
||||||
|
admin_pool_provider_id_from_scores_path, build_admin_pool_error_response,
|
||||||
|
parse_admin_pool_page, parse_admin_pool_page_size,
|
||||||
|
};
|
||||||
|
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
|
||||||
|
use crate::handlers::shared::query_param_value;
|
||||||
|
use crate::GatewayError;
|
||||||
|
use aether_data_contracts::repository::pool_scores::{
|
||||||
|
ListPoolMemberScoresQuery, PoolMemberHardState, PoolMemberProbeStatus,
|
||||||
|
POOL_KIND_PROVIDER_KEY_POOL, POOL_SCORE_SCOPE_KIND_MODEL,
|
||||||
|
};
|
||||||
|
use axum::{
|
||||||
|
body::Body,
|
||||||
|
http,
|
||||||
|
response::{IntoResponse, Response},
|
||||||
|
Json,
|
||||||
|
};
|
||||||
|
use serde_json::json;
|
||||||
|
use std::collections::BTreeMap;
|
||||||
|
|
||||||
|
pub(super) async fn build_admin_pool_scores_response(
|
||||||
|
state: &AdminAppState<'_>,
|
||||||
|
request_context: &AdminRequestContext<'_>,
|
||||||
|
) -> Result<Response<Body>, GatewayError> {
|
||||||
|
let Some(provider_id) = admin_pool_provider_id_from_scores_path(request_context.path()) else {
|
||||||
|
return Ok(build_admin_pool_error_response(
|
||||||
|
http::StatusCode::BAD_REQUEST,
|
||||||
|
"provider_id 无效",
|
||||||
|
));
|
||||||
|
};
|
||||||
|
|
||||||
|
let query = request_context.query_string();
|
||||||
|
let page = match parse_admin_pool_page(query) {
|
||||||
|
Ok(value) => value,
|
||||||
|
Err(message) => {
|
||||||
|
return Ok(build_admin_pool_error_response(
|
||||||
|
http::StatusCode::BAD_REQUEST,
|
||||||
|
message,
|
||||||
|
));
|
||||||
|
}
|
||||||
|
};
|
||||||
|
let page_size = match parse_admin_pool_page_size(query) {
|
||||||
|
Ok(value) => value.min(500),
|
||||||
|
Err(message) => {
|
||||||
|
return Ok(build_admin_pool_error_response(
|
||||||
|
http::StatusCode::BAD_REQUEST,
|
||||||
|
message,
|
||||||
|
));
|
||||||
|
}
|
||||||
|
};
|
||||||
|
let offset = page.saturating_sub(1).saturating_mul(page_size);
|
||||||
|
let api_format = query_param_value(query, "api_format")
|
||||||
|
.map(|value| crate::ai_serving::normalize_api_format_alias(value.as_str()));
|
||||||
|
let model_id = query_param_value(query, "model_id");
|
||||||
|
let hard_states = match parse_hard_state_filter(query) {
|
||||||
|
Ok(value) => value,
|
||||||
|
Err(message) => {
|
||||||
|
return Ok(build_admin_pool_error_response(
|
||||||
|
http::StatusCode::BAD_REQUEST,
|
||||||
|
message,
|
||||||
|
));
|
||||||
|
}
|
||||||
|
};
|
||||||
|
let probe_statuses = match parse_probe_status_filter(query) {
|
||||||
|
Ok(value) => value,
|
||||||
|
Err(message) => {
|
||||||
|
return Ok(build_admin_pool_error_response(
|
||||||
|
http::StatusCode::BAD_REQUEST,
|
||||||
|
message,
|
||||||
|
));
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
let scores = state
|
||||||
|
.app()
|
||||||
|
.data
|
||||||
|
.list_pool_member_scores(&ListPoolMemberScoresQuery {
|
||||||
|
pool_kind: POOL_KIND_PROVIDER_KEY_POOL.to_string(),
|
||||||
|
pool_id: provider_id.clone(),
|
||||||
|
capability: api_format.clone(),
|
||||||
|
scope_kind: Some(POOL_SCORE_SCOPE_KIND_MODEL.to_string()),
|
||||||
|
scope_id: model_id.clone(),
|
||||||
|
hard_states,
|
||||||
|
probe_statuses,
|
||||||
|
offset,
|
||||||
|
limit: page_size,
|
||||||
|
})
|
||||||
|
.await
|
||||||
|
.map_err(|err| GatewayError::Internal(format!("{err:?}")))?;
|
||||||
|
|
||||||
|
let key_ids = scores
|
||||||
|
.iter()
|
||||||
|
.map(|score| score.member_id.clone())
|
||||||
|
.collect::<Vec<_>>();
|
||||||
|
let keys = state
|
||||||
|
.app()
|
||||||
|
.read_provider_catalog_keys_by_ids(&key_ids)
|
||||||
|
.await
|
||||||
|
.unwrap_or_default()
|
||||||
|
.into_iter()
|
||||||
|
.map(|key| (key.id.clone(), key))
|
||||||
|
.collect::<BTreeMap<_, _>>();
|
||||||
|
|
||||||
|
let items = scores
|
||||||
|
.into_iter()
|
||||||
|
.map(|score| {
|
||||||
|
let key = keys.get(&score.member_id);
|
||||||
|
json!({
|
||||||
|
"id": score.id,
|
||||||
|
"pool_kind": score.pool_kind,
|
||||||
|
"pool_id": score.pool_id,
|
||||||
|
"member_kind": score.member_kind,
|
||||||
|
"member_id": score.member_id,
|
||||||
|
"capability": score.capability,
|
||||||
|
"scope_kind": score.scope_kind,
|
||||||
|
"scope_id": score.scope_id,
|
||||||
|
"score": score.score,
|
||||||
|
"hard_state": score.hard_state.as_database(),
|
||||||
|
"score_version": score.score_version,
|
||||||
|
"score_reason": score.score_reason,
|
||||||
|
"last_ranked_at": score.last_ranked_at,
|
||||||
|
"last_scheduled_at": score.last_scheduled_at,
|
||||||
|
"last_success_at": score.last_success_at,
|
||||||
|
"last_failure_at": score.last_failure_at,
|
||||||
|
"failure_count": score.failure_count,
|
||||||
|
"last_probe_attempt_at": score.last_probe_attempt_at,
|
||||||
|
"last_probe_success_at": score.last_probe_success_at,
|
||||||
|
"last_probe_failure_at": score.last_probe_failure_at,
|
||||||
|
"probe_failure_count": score.probe_failure_count,
|
||||||
|
"probe_status": score.probe_status.as_database(),
|
||||||
|
"updated_at": score.updated_at,
|
||||||
|
"key": key.map(|key| json!({
|
||||||
|
"id": key.id,
|
||||||
|
"name": key.name,
|
||||||
|
"auth_type": key.auth_type,
|
||||||
|
"is_active": key.is_active,
|
||||||
|
"internal_priority": key.internal_priority,
|
||||||
|
"last_used_at": key.last_used_at_unix_secs,
|
||||||
|
}))
|
||||||
|
})
|
||||||
|
})
|
||||||
|
.collect::<Vec<_>>();
|
||||||
|
|
||||||
|
Ok(Json(json!({
|
||||||
|
"provider_id": provider_id,
|
||||||
|
"page": page,
|
||||||
|
"page_size": page_size,
|
||||||
|
"filters": {
|
||||||
|
"api_format": api_format,
|
||||||
|
"model_id": model_id,
|
||||||
|
"hard_state": query_param_value(query, "hard_state"),
|
||||||
|
"probe_status": query_param_value(query, "probe_status")
|
||||||
|
},
|
||||||
|
"items": items
|
||||||
|
}))
|
||||||
|
.into_response())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn parse_hard_state_filter(query: Option<&str>) -> Result<Vec<PoolMemberHardState>, String> {
|
||||||
|
let Some(raw) = query_param_value(query, "hard_state") else {
|
||||||
|
return Ok(Vec::new());
|
||||||
|
};
|
||||||
|
raw.split(',')
|
||||||
|
.map(|value| match value.trim() {
|
||||||
|
"available" => Ok(PoolMemberHardState::Available),
|
||||||
|
"unknown" => Ok(PoolMemberHardState::Unknown),
|
||||||
|
"cooldown" => Ok(PoolMemberHardState::Cooldown),
|
||||||
|
"quota_exhausted" => Ok(PoolMemberHardState::QuotaExhausted),
|
||||||
|
"auth_invalid" => Ok(PoolMemberHardState::AuthInvalid),
|
||||||
|
"banned" => Ok(PoolMemberHardState::Banned),
|
||||||
|
"inactive" => Ok(PoolMemberHardState::Inactive),
|
||||||
|
_ => Err("hard_state must be one of: available, unknown, cooldown, quota_exhausted, auth_invalid, banned, inactive".to_string()),
|
||||||
|
})
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn parse_probe_status_filter(
|
||||||
|
query: Option<&str>,
|
||||||
|
) -> Result<Option<Vec<PoolMemberProbeStatus>>, String> {
|
||||||
|
let Some(raw) = query_param_value(query, "probe_status") else {
|
||||||
|
return Ok(None);
|
||||||
|
};
|
||||||
|
raw.split(',')
|
||||||
|
.map(|value| match value.trim() {
|
||||||
|
"never" => Ok(PoolMemberProbeStatus::Never),
|
||||||
|
"ok" => Ok(PoolMemberProbeStatus::Ok),
|
||||||
|
"failed" => Ok(PoolMemberProbeStatus::Failed),
|
||||||
|
"stale" => Ok(PoolMemberProbeStatus::Stale),
|
||||||
|
"in_progress" => Ok(PoolMemberProbeStatus::InProgress),
|
||||||
|
_ => Err(
|
||||||
|
"probe_status must be one of: never, ok, failed, stale, in_progress".to_string(),
|
||||||
|
),
|
||||||
|
})
|
||||||
|
.collect::<Result<Vec<_>, _>>()
|
||||||
|
.map(Some)
|
||||||
|
}
|
||||||
@@ -149,6 +149,18 @@ pub(crate) fn admin_pool_provider_id_from_path(request_path: &str) -> Option<Str
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub(crate) fn admin_pool_provider_id_from_scores_path(request_path: &str) -> Option<String> {
|
||||||
|
let raw = request_path.strip_prefix("/api/admin/pool/")?;
|
||||||
|
let mut segments = raw.split('/');
|
||||||
|
let provider_id = segments.next()?.trim();
|
||||||
|
let scores_segment = segments.next()?.trim_end_matches('/').trim();
|
||||||
|
if provider_id.is_empty() || scores_segment != "scores" {
|
||||||
|
None
|
||||||
|
} else {
|
||||||
|
Some(provider_id.to_string())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
pub(crate) fn is_admin_pool_route(request_context: &AdminRequestContext<'_>) -> bool {
|
pub(crate) fn is_admin_pool_route(request_context: &AdminRequestContext<'_>) -> bool {
|
||||||
let normalized_path = request_context.path().trim_end_matches('/');
|
let normalized_path = request_context.path().trim_end_matches('/');
|
||||||
let path = if normalized_path.is_empty() {
|
let path = if normalized_path.is_empty() {
|
||||||
@@ -164,6 +176,10 @@ pub(crate) fn is_admin_pool_route(request_context: &AdminRequestContext<'_>) ->
|
|||||||
&& path.starts_with("/api/admin/pool/")
|
&& path.starts_with("/api/admin/pool/")
|
||||||
&& path.ends_with("/keys")
|
&& path.ends_with("/keys")
|
||||||
&& path.matches('/').count() == 5)
|
&& path.matches('/').count() == 5)
|
||||||
|
|| (request_context.method() == http::Method::GET
|
||||||
|
&& path.starts_with("/api/admin/pool/")
|
||||||
|
&& path.ends_with("/scores")
|
||||||
|
&& path.matches('/').count() == 5)
|
||||||
|| (request_context.method() == http::Method::POST
|
|| (request_context.method() == http::Method::POST
|
||||||
&& path.starts_with("/api/admin/pool/")
|
&& path.starts_with("/api/admin/pool/")
|
||||||
&& path.ends_with("/keys/batch-import")
|
&& path.ends_with("/keys/batch-import")
|
||||||
|
|||||||
@@ -1,5 +1,6 @@
|
|||||||
use crate::handlers::admin::request::AdminAppState;
|
use crate::handlers::admin::request::AdminAppState;
|
||||||
use crate::LocalProviderDeleteTaskState;
|
use crate::LocalProviderDeleteTaskState;
|
||||||
|
use aether_ai_serving::PoolMemberScoreRules;
|
||||||
use serde_json::json;
|
use serde_json::json;
|
||||||
use std::collections::BTreeMap;
|
use std::collections::BTreeMap;
|
||||||
|
|
||||||
@@ -39,6 +40,10 @@ pub(crate) struct AdminProviderPoolConfig {
|
|||||||
pub(crate) health_policy_enabled: bool,
|
pub(crate) health_policy_enabled: bool,
|
||||||
pub(crate) probing_enabled: bool,
|
pub(crate) probing_enabled: bool,
|
||||||
pub(crate) probing_interval_minutes: u64,
|
pub(crate) probing_interval_minutes: u64,
|
||||||
|
pub(crate) probe_concurrency: u64,
|
||||||
|
pub(crate) score_top_n: u64,
|
||||||
|
pub(crate) score_fallback_scan_limit: u64,
|
||||||
|
pub(crate) score_rules: PoolMemberScoreRules,
|
||||||
pub(crate) stream_timeout_threshold: u64,
|
pub(crate) stream_timeout_threshold: u64,
|
||||||
pub(crate) stream_timeout_window_seconds: u64,
|
pub(crate) stream_timeout_window_seconds: u64,
|
||||||
pub(crate) stream_timeout_cooldown_seconds: u64,
|
pub(crate) stream_timeout_cooldown_seconds: u64,
|
||||||
|
|||||||
@@ -4,14 +4,15 @@ mod tests;
|
|||||||
|
|
||||||
pub(crate) use runtime::{
|
pub(crate) use runtime::{
|
||||||
cancel_proxy_upgrade_rollout, clear_proxy_upgrade_rollout_conflicts,
|
cancel_proxy_upgrade_rollout, clear_proxy_upgrade_rollout_conflicts,
|
||||||
inspect_proxy_upgrade_rollout, list_admin_cleanup_run_records,
|
ensure_provider_key_pool_scores_for_keys, inspect_proxy_upgrade_rollout,
|
||||||
perform_oauth_token_refresh_once, perform_pool_quota_probe_once, perform_provider_checkin_once,
|
list_admin_cleanup_run_records, perform_oauth_token_refresh_once,
|
||||||
rebuild_admin_stats_once, record_completed_cleanup_run, record_proxy_upgrade_traffic_success,
|
perform_pool_quota_probe_once, perform_provider_checkin_once, rebuild_admin_stats_once,
|
||||||
|
record_completed_cleanup_run, record_proxy_upgrade_traffic_success,
|
||||||
restore_proxy_upgrade_rollout_skipped_nodes, retry_proxy_upgrade_rollout_node,
|
restore_proxy_upgrade_rollout_skipped_nodes, retry_proxy_upgrade_rollout_node,
|
||||||
run_admin_system_cleanup_once, skip_proxy_upgrade_rollout_node, spawn_audit_cleanup_worker,
|
run_admin_system_cleanup_once, skip_proxy_upgrade_rollout_node, spawn_audit_cleanup_worker,
|
||||||
spawn_db_maintenance_worker, spawn_gemini_file_mapping_cleanup_worker,
|
spawn_db_maintenance_worker, spawn_gemini_file_mapping_cleanup_worker,
|
||||||
spawn_oauth_token_refresh_worker, spawn_pending_cleanup_worker, spawn_pool_monitor_worker,
|
spawn_oauth_token_refresh_worker, spawn_pending_cleanup_worker, spawn_pool_monitor_worker,
|
||||||
spawn_pool_quota_probe_worker, spawn_provider_checkin_worker,
|
spawn_pool_quota_probe_worker, spawn_pool_score_rebuild_worker, spawn_provider_checkin_worker,
|
||||||
spawn_proxy_node_metrics_cleanup_worker, spawn_proxy_node_stale_cleanup_worker,
|
spawn_proxy_node_metrics_cleanup_worker, spawn_proxy_node_stale_cleanup_worker,
|
||||||
spawn_proxy_upgrade_rollout_worker, spawn_request_candidate_cleanup_worker,
|
spawn_proxy_upgrade_rollout_worker, spawn_request_candidate_cleanup_worker,
|
||||||
spawn_stats_aggregation_worker, spawn_stats_hourly_aggregation_worker,
|
spawn_stats_aggregation_worker, spawn_stats_hourly_aggregation_worker,
|
||||||
|
|||||||
@@ -20,6 +20,8 @@ mod oauth_token_refresh;
|
|||||||
mod pending_cleanup;
|
mod pending_cleanup;
|
||||||
#[path = "runtime/pool_quota_probe.rs"]
|
#[path = "runtime/pool_quota_probe.rs"]
|
||||||
mod pool_quota_probe;
|
mod pool_quota_probe;
|
||||||
|
#[path = "runtime/pool_score_rebuild.rs"]
|
||||||
|
mod pool_score_rebuild;
|
||||||
#[path = "runtime/provider_checkin.rs"]
|
#[path = "runtime/provider_checkin.rs"]
|
||||||
mod provider_checkin;
|
mod provider_checkin;
|
||||||
#[path = "runtime/proxy_node_metrics_cleanup.rs"]
|
#[path = "runtime/proxy_node_metrics_cleanup.rs"]
|
||||||
@@ -67,6 +69,11 @@ pub(crate) use pool_quota_probe::{
|
|||||||
select_pool_quota_probe_key_ids, spawn_pool_quota_probe_worker, PoolQuotaProbeRunSummary,
|
select_pool_quota_probe_key_ids, spawn_pool_quota_probe_worker, PoolQuotaProbeRunSummary,
|
||||||
PoolQuotaProbeWorkerConfig,
|
PoolQuotaProbeWorkerConfig,
|
||||||
};
|
};
|
||||||
|
pub(crate) use pool_score_rebuild::{
|
||||||
|
ensure_provider_key_pool_scores_for_keys, perform_pool_score_rebuild_once,
|
||||||
|
perform_pool_score_rebuild_once_with_config, spawn_pool_score_rebuild_worker,
|
||||||
|
PoolScoreRebuildRunSummary, PoolScoreRebuildWorkerConfig,
|
||||||
|
};
|
||||||
pub(crate) use provider_checkin::{perform_provider_checkin_once, ProviderCheckinRunSummary};
|
pub(crate) use provider_checkin::{perform_provider_checkin_once, ProviderCheckinRunSummary};
|
||||||
use proxy_node_metrics_cleanup::*;
|
use proxy_node_metrics_cleanup::*;
|
||||||
use proxy_node_staleness::*;
|
use proxy_node_staleness::*;
|
||||||
|
|||||||
@@ -1,10 +1,16 @@
|
|||||||
use std::collections::BTreeMap;
|
use std::collections::BTreeMap;
|
||||||
use std::time::{Duration, SystemTime, UNIX_EPOCH};
|
use std::time::{Duration, SystemTime, UNIX_EPOCH};
|
||||||
|
|
||||||
|
use aether_data_contracts::repository::pool_scores::{
|
||||||
|
ListPoolMemberProbeCandidatesQuery, PoolMemberHardState, PoolMemberIdentity,
|
||||||
|
PoolMemberProbeAttempt, PoolMemberProbeResult, PoolMemberProbeStatus,
|
||||||
|
POOL_KIND_PROVIDER_KEY_POOL,
|
||||||
|
};
|
||||||
use aether_data_contracts::repository::provider_catalog::{
|
use aether_data_contracts::repository::provider_catalog::{
|
||||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
|
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
|
||||||
};
|
};
|
||||||
use aether_runtime_state::{RuntimeLockLease, RuntimeState};
|
use aether_runtime_state::{RuntimeLockLease, RuntimeState};
|
||||||
|
use futures_util::{stream, StreamExt};
|
||||||
use serde_json::Value;
|
use serde_json::Value;
|
||||||
use tracing::{debug, info, warn};
|
use tracing::{debug, info, warn};
|
||||||
|
|
||||||
@@ -15,10 +21,13 @@ use crate::admin_api::{
|
|||||||
};
|
};
|
||||||
use crate::{AppState, GatewayError};
|
use crate::{AppState, GatewayError};
|
||||||
|
|
||||||
|
use super::pool_score_rebuild::ensure_provider_key_pool_scores_for_keys;
|
||||||
|
|
||||||
const POOL_QUOTA_PROBE_REDIS_PREFIX: &str = "ap:quota_probe:last";
|
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_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;
|
||||||
|
|
||||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||||
@@ -50,6 +59,7 @@ impl PoolQuotaProbeRunSummary {
|
|||||||
pub(crate) struct PoolQuotaProbeWorkerConfig {
|
pub(crate) struct PoolQuotaProbeWorkerConfig {
|
||||||
pub(crate) scan_interval: Duration,
|
pub(crate) scan_interval: Duration,
|
||||||
pub(crate) max_keys_per_provider: usize,
|
pub(crate) max_keys_per_provider: usize,
|
||||||
|
pub(crate) global_concurrency: usize,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl PoolQuotaProbeWorkerConfig {
|
impl PoolQuotaProbeWorkerConfig {
|
||||||
@@ -63,9 +73,15 @@ impl PoolQuotaProbeWorkerConfig {
|
|||||||
"POOL_QUOTA_PROBE_MAX_KEYS_PER_PROVIDER",
|
"POOL_QUOTA_PROBE_MAX_KEYS_PER_PROVIDER",
|
||||||
POOL_QUOTA_PROBE_DEFAULT_MAX_KEYS_PER_PROVIDER,
|
POOL_QUOTA_PROBE_DEFAULT_MAX_KEYS_PER_PROVIDER,
|
||||||
);
|
);
|
||||||
|
let global_concurrency = env_usize(
|
||||||
|
"POOL_QUOTA_PROBE_GLOBAL_CONCURRENCY",
|
||||||
|
POOL_QUOTA_PROBE_DEFAULT_GLOBAL_CONCURRENCY,
|
||||||
|
)
|
||||||
|
.clamp(1, 256);
|
||||||
Self {
|
Self {
|
||||||
scan_interval: Duration::from_secs(scan_interval_seconds),
|
scan_interval: Duration::from_secs(scan_interval_seconds),
|
||||||
max_keys_per_provider,
|
max_keys_per_provider,
|
||||||
|
global_concurrency,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -172,6 +188,49 @@ pub(crate) fn select_pool_quota_probe_key_ids(
|
|||||||
stale.into_iter().map(|(_, key_id)| key_id).collect()
|
stale.into_iter().map(|(_, key_id)| key_id).collect()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
async fn select_score_probe_key_ids(
|
||||||
|
state: &AppState,
|
||||||
|
provider_id: &str,
|
||||||
|
now_ts: u64,
|
||||||
|
interval_seconds: u64,
|
||||||
|
limit: usize,
|
||||||
|
) -> Vec<String> {
|
||||||
|
if limit == 0 {
|
||||||
|
return Vec::new();
|
||||||
|
}
|
||||||
|
let stale_before_unix_secs = now_ts.saturating_sub(interval_seconds);
|
||||||
|
let query = ListPoolMemberProbeCandidatesQuery {
|
||||||
|
pool_kind: POOL_KIND_PROVIDER_KEY_POOL.to_string(),
|
||||||
|
pool_id: provider_id.to_string(),
|
||||||
|
capability: None,
|
||||||
|
stale_before_unix_secs,
|
||||||
|
limit: limit.saturating_mul(4).max(limit),
|
||||||
|
};
|
||||||
|
let scores = match state.data.list_pool_member_probe_candidates(&query).await {
|
||||||
|
Ok(scores) => scores,
|
||||||
|
Err(err) => {
|
||||||
|
debug!(
|
||||||
|
provider_id,
|
||||||
|
error = ?err,
|
||||||
|
"gateway pool quota probe: failed to read score probe candidates"
|
||||||
|
);
|
||||||
|
return Vec::new();
|
||||||
|
}
|
||||||
|
};
|
||||||
|
let mut selected = Vec::new();
|
||||||
|
let mut seen = std::collections::BTreeSet::new();
|
||||||
|
for score in scores {
|
||||||
|
if !seen.insert(score.member_id.clone()) {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
selected.push(score.member_id);
|
||||||
|
if selected.len() >= limit {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
selected
|
||||||
|
}
|
||||||
|
|
||||||
fn probe_stamp_key(provider_id: &str, key_id: &str) -> String {
|
fn probe_stamp_key(provider_id: &str, key_id: &str) -> String {
|
||||||
format!("{POOL_QUOTA_PROBE_REDIS_PREFIX}:{provider_id}:{key_id}")
|
format!("{POOL_QUOTA_PROBE_REDIS_PREFIX}:{provider_id}:{key_id}")
|
||||||
}
|
}
|
||||||
@@ -295,7 +354,28 @@ async fn select_keys_for_provider(
|
|||||||
|
|
||||||
let key_ids = keys.iter().map(|key| key.id.clone()).collect::<Vec<_>>();
|
let key_ids = keys.iter().map(|key| key.id.clone()).collect::<Vec<_>>();
|
||||||
let probe_stamps = load_probe_timestamps(runtime, &provider.id, &key_ids).await;
|
let probe_stamps = load_probe_timestamps(runtime, &provider.id, &key_ids).await;
|
||||||
let selected_ids = select_pool_quota_probe_key_ids(
|
let mut selected_ids = select_score_probe_key_ids(
|
||||||
|
state,
|
||||||
|
&provider.id,
|
||||||
|
now_ts,
|
||||||
|
interval_seconds,
|
||||||
|
max_keys_per_provider,
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.into_iter()
|
||||||
|
.filter(|key_id| key_ids.iter().any(|known_id| known_id == key_id))
|
||||||
|
.filter(|key_id| {
|
||||||
|
probe_stamps.get(key_id).is_none_or(|last_probe_ts| {
|
||||||
|
now_ts.saturating_sub(*last_probe_ts) >= interval_seconds
|
||||||
|
})
|
||||||
|
})
|
||||||
|
.collect::<Vec<_>>();
|
||||||
|
let mut selected_seen = selected_ids
|
||||||
|
.iter()
|
||||||
|
.cloned()
|
||||||
|
.collect::<std::collections::BTreeSet<_>>();
|
||||||
|
let remaining = max_keys_per_provider.saturating_sub(selected_ids.len());
|
||||||
|
let fallback_selected_ids = select_pool_quota_probe_key_ids(
|
||||||
&keys,
|
&keys,
|
||||||
provider_type,
|
provider_type,
|
||||||
now_ts,
|
now_ts,
|
||||||
@@ -303,6 +383,14 @@ async fn select_keys_for_provider(
|
|||||||
&probe_stamps,
|
&probe_stamps,
|
||||||
max_keys_per_provider,
|
max_keys_per_provider,
|
||||||
);
|
);
|
||||||
|
for key_id in fallback_selected_ids {
|
||||||
|
if selected_ids.len() >= max_keys_per_provider || remaining == 0 {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
if selected_seen.insert(key_id.clone()) {
|
||||||
|
selected_ids.push(key_id);
|
||||||
|
}
|
||||||
|
}
|
||||||
if selected_ids.is_empty() {
|
if selected_ids.is_empty() {
|
||||||
return Ok(Vec::new());
|
return Ok(Vec::new());
|
||||||
}
|
}
|
||||||
@@ -381,6 +469,168 @@ fn update_summary_from_payload(
|
|||||||
.unwrap_or(0) as usize;
|
.unwrap_or(0) as usize;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
async fn record_score_probe_results_from_payload(
|
||||||
|
state: &AppState,
|
||||||
|
provider_id: &str,
|
||||||
|
selected_key_ids: &[String],
|
||||||
|
payload: Option<&Value>,
|
||||||
|
attempted_at: u64,
|
||||||
|
) {
|
||||||
|
let mut recorded = std::collections::BTreeSet::new();
|
||||||
|
if let Some(results) = payload
|
||||||
|
.and_then(|value| value.get("results"))
|
||||||
|
.and_then(Value::as_array)
|
||||||
|
{
|
||||||
|
for item in results {
|
||||||
|
let Some(key_id) = item
|
||||||
|
.get("key_id")
|
||||||
|
.and_then(Value::as_str)
|
||||||
|
.map(str::trim)
|
||||||
|
.filter(|value| !value.is_empty())
|
||||||
|
else {
|
||||||
|
continue;
|
||||||
|
};
|
||||||
|
recorded.insert(key_id.to_string());
|
||||||
|
record_score_probe_result_for_key(
|
||||||
|
state,
|
||||||
|
provider_id,
|
||||||
|
key_id,
|
||||||
|
attempted_at,
|
||||||
|
probe_result_succeeded(item),
|
||||||
|
probe_result_hard_state(item).or_else(|| {
|
||||||
|
(!probe_result_succeeded(item)).then_some(PoolMemberHardState::Cooldown)
|
||||||
|
}),
|
||||||
|
serde_json::json!({
|
||||||
|
"last_probe": {
|
||||||
|
"source": "pool_quota_probe",
|
||||||
|
"status": item.get("status").cloned().unwrap_or(Value::Null),
|
||||||
|
"status_code": item.get("status_code").cloned().unwrap_or(Value::Null),
|
||||||
|
"message": item.get("message").cloned().unwrap_or(Value::Null),
|
||||||
|
"auto_removed": item.get("auto_removed").cloned().unwrap_or(Value::Null)
|
||||||
|
}
|
||||||
|
}),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
for key_id in selected_key_ids {
|
||||||
|
if recorded.contains(key_id) {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
record_score_probe_result_for_key(
|
||||||
|
state,
|
||||||
|
provider_id,
|
||||||
|
key_id,
|
||||||
|
attempted_at,
|
||||||
|
false,
|
||||||
|
Some(PoolMemberHardState::Cooldown),
|
||||||
|
serde_json::json!({
|
||||||
|
"last_probe": {
|
||||||
|
"source": "pool_quota_probe",
|
||||||
|
"status": "missing_result"
|
||||||
|
}
|
||||||
|
}),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn record_score_probe_result_for_key(
|
||||||
|
state: &AppState,
|
||||||
|
provider_id: &str,
|
||||||
|
key_id: &str,
|
||||||
|
attempted_at: u64,
|
||||||
|
succeeded: bool,
|
||||||
|
hard_state: Option<PoolMemberHardState>,
|
||||||
|
score_reason_patch: Value,
|
||||||
|
) {
|
||||||
|
let result = PoolMemberProbeResult {
|
||||||
|
identity: PoolMemberIdentity::provider_api_key(provider_id.to_string(), key_id.to_string()),
|
||||||
|
scope: None,
|
||||||
|
attempted_at,
|
||||||
|
succeeded,
|
||||||
|
hard_state,
|
||||||
|
probe_status: if succeeded {
|
||||||
|
PoolMemberProbeStatus::Ok
|
||||||
|
} else {
|
||||||
|
PoolMemberProbeStatus::Failed
|
||||||
|
},
|
||||||
|
score_reason_patch: Some(score_reason_patch),
|
||||||
|
};
|
||||||
|
if let Err(err) = state.data.record_pool_member_probe_result(result).await {
|
||||||
|
debug!(
|
||||||
|
provider_id,
|
||||||
|
key_id,
|
||||||
|
error = ?err,
|
||||||
|
"gateway pool quota probe: failed to record score probe result"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn record_score_probe_in_progress_for_key(
|
||||||
|
state: &AppState,
|
||||||
|
provider_id: &str,
|
||||||
|
key_id: &str,
|
||||||
|
attempted_at: u64,
|
||||||
|
) {
|
||||||
|
let attempt = PoolMemberProbeAttempt {
|
||||||
|
identity: PoolMemberIdentity::provider_api_key(provider_id.to_string(), key_id.to_string()),
|
||||||
|
scope: None,
|
||||||
|
attempted_at,
|
||||||
|
score_reason_patch: Some(serde_json::json!({
|
||||||
|
"last_probe": {
|
||||||
|
"source": "pool_quota_probe",
|
||||||
|
"status": "in_progress"
|
||||||
|
}
|
||||||
|
})),
|
||||||
|
};
|
||||||
|
if let Err(err) = state.data.mark_pool_member_probe_in_progress(attempt).await {
|
||||||
|
debug!(
|
||||||
|
provider_id,
|
||||||
|
key_id,
|
||||||
|
error = ?err,
|
||||||
|
"gateway pool quota probe: failed to mark score probe in progress"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn probe_result_succeeded(item: &Value) -> bool {
|
||||||
|
item.get("status")
|
||||||
|
.and_then(Value::as_str)
|
||||||
|
.is_some_and(|status| status == "success")
|
||||||
|
}
|
||||||
|
|
||||||
|
fn probe_result_hard_state(item: &Value) -> Option<PoolMemberHardState> {
|
||||||
|
if probe_result_succeeded(item) {
|
||||||
|
return Some(PoolMemberHardState::Available);
|
||||||
|
}
|
||||||
|
if item
|
||||||
|
.get("auto_removed")
|
||||||
|
.and_then(Value::as_bool)
|
||||||
|
.unwrap_or(false)
|
||||||
|
{
|
||||||
|
return Some(PoolMemberHardState::Banned);
|
||||||
|
}
|
||||||
|
let status = item
|
||||||
|
.get("status")
|
||||||
|
.and_then(Value::as_str)
|
||||||
|
.unwrap_or_default()
|
||||||
|
.trim()
|
||||||
|
.to_ascii_lowercase();
|
||||||
|
match status.as_str() {
|
||||||
|
"auth_invalid" | "forbidden" => Some(PoolMemberHardState::AuthInvalid),
|
||||||
|
"workspace_deactivated" => Some(PoolMemberHardState::Banned),
|
||||||
|
"quota_exhausted" => Some(PoolMemberHardState::QuotaExhausted),
|
||||||
|
_ => match item.get("status_code").and_then(Value::as_u64) {
|
||||||
|
Some(401 | 403) => Some(PoolMemberHardState::AuthInvalid),
|
||||||
|
Some(402) => Some(PoolMemberHardState::QuotaExhausted),
|
||||||
|
Some(429 | 500..=599) => Some(PoolMemberHardState::Cooldown),
|
||||||
|
_ => None,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
pub(crate) async fn perform_pool_quota_probe_once_with_config(
|
pub(crate) async fn perform_pool_quota_probe_once_with_config(
|
||||||
state: &AppState,
|
state: &AppState,
|
||||||
config: PoolQuotaProbeWorkerConfig,
|
config: PoolQuotaProbeWorkerConfig,
|
||||||
@@ -404,11 +654,7 @@ pub(crate) async fn perform_pool_quota_probe_once_with_config(
|
|||||||
.filter_map(|(provider, provider_type)| {
|
.filter_map(|(provider, provider_type)| {
|
||||||
let pool_config = admin_provider_pool_config(&provider)?;
|
let pool_config = admin_provider_pool_config(&provider)?;
|
||||||
if pool_config.probing_enabled {
|
if pool_config.probing_enabled {
|
||||||
Some((
|
Some((provider, provider_type, pool_config))
|
||||||
provider,
|
|
||||||
provider_type,
|
|
||||||
pool_config.probing_interval_minutes,
|
|
||||||
))
|
|
||||||
} else {
|
} else {
|
||||||
None
|
None
|
||||||
}
|
}
|
||||||
@@ -441,7 +687,7 @@ pub(crate) async fn perform_pool_quota_probe_once_with_config(
|
|||||||
..PoolQuotaProbeRunSummary::empty()
|
..PoolQuotaProbeRunSummary::empty()
|
||||||
};
|
};
|
||||||
|
|
||||||
for (provider, provider_type, interval_minutes) in providers {
|
for (provider, provider_type, pool_config) in providers {
|
||||||
let endpoints = endpoints_by_provider
|
let endpoints = endpoints_by_provider
|
||||||
.remove(&provider.id)
|
.remove(&provider.id)
|
||||||
.unwrap_or_default();
|
.unwrap_or_default();
|
||||||
@@ -455,6 +701,7 @@ pub(crate) async fn perform_pool_quota_probe_once_with_config(
|
|||||||
continue;
|
continue;
|
||||||
};
|
};
|
||||||
|
|
||||||
|
let interval_minutes = pool_config.probing_interval_minutes;
|
||||||
let interval_seconds = interval_minutes.clamp(1, 1440).saturating_mul(60);
|
let interval_seconds = interval_minutes.clamp(1, 1440).saturating_mul(60);
|
||||||
let keys = select_keys_for_provider(
|
let keys = select_keys_for_provider(
|
||||||
state,
|
state,
|
||||||
@@ -474,42 +721,131 @@ pub(crate) async fn perform_pool_quota_probe_once_with_config(
|
|||||||
summary.providers_probed += 1;
|
summary.providers_probed += 1;
|
||||||
summary.selected_keys += selected_count;
|
summary.selected_keys += selected_count;
|
||||||
|
|
||||||
let provider_short_id = provider.id.chars().take(8).collect::<String>();
|
let selected_key_ids = keys.iter().map(|key| key.id.clone()).collect::<Vec<_>>();
|
||||||
match refresh_provider_probe_keys(&admin_state, &provider, &endpoint, &provider_type, keys)
|
let score_ensure_budget = (pool_config.score_fallback_scan_limit as usize)
|
||||||
.await
|
.min(50_000)
|
||||||
|
.max(selected_count.min(50_000));
|
||||||
|
match ensure_provider_key_pool_scores_for_keys(
|
||||||
|
state,
|
||||||
|
&provider,
|
||||||
|
&pool_config,
|
||||||
|
&endpoints,
|
||||||
|
&keys,
|
||||||
|
now_ts,
|
||||||
|
score_ensure_budget,
|
||||||
|
)
|
||||||
|
.await
|
||||||
{
|
{
|
||||||
Ok(payload) => {
|
Ok(upserted) if upserted > 0 => {
|
||||||
update_summary_from_payload(&mut summary, selected_count, payload.as_ref());
|
debug!(
|
||||||
let probe_success = payload
|
provider_id = %provider.id,
|
||||||
.as_ref()
|
key_count = selected_count,
|
||||||
.and_then(|value| value.get("success"))
|
scores_upserted = upserted,
|
||||||
.and_then(Value::as_u64)
|
"gateway pool quota probe: ensured score rows for selected probe keys"
|
||||||
.unwrap_or(0);
|
|
||||||
let probe_failed = payload
|
|
||||||
.as_ref()
|
|
||||||
.and_then(|value| value.get("failed"))
|
|
||||||
.and_then(Value::as_u64)
|
|
||||||
.unwrap_or(0);
|
|
||||||
info!(
|
|
||||||
provider_id = %provider_short_id,
|
|
||||||
provider_type,
|
|
||||||
selected = selected_count,
|
|
||||||
success = probe_success,
|
|
||||||
failed = probe_failed,
|
|
||||||
"gateway pool quota probe completed"
|
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
Ok(_) => {}
|
||||||
Err(err) => {
|
Err(err) => {
|
||||||
summary.failed += selected_count;
|
|
||||||
warn!(
|
warn!(
|
||||||
provider_id = %provider_short_id,
|
provider_id = %provider.id,
|
||||||
provider_type,
|
key_count = selected_count,
|
||||||
selected = selected_count,
|
|
||||||
error = ?err,
|
error = ?err,
|
||||||
"gateway pool quota probe failed"
|
"gateway pool quota probe: failed to ensure score rows for selected probe keys"
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
for key_id in &selected_key_ids {
|
||||||
|
record_score_probe_in_progress_for_key(state, &provider.id, key_id, now_ts).await;
|
||||||
|
}
|
||||||
|
|
||||||
|
let provider_short_id = provider.id.chars().take(8).collect::<String>();
|
||||||
|
let probe_concurrency = pool_config.probe_concurrency.clamp(1, 64) as usize;
|
||||||
|
let probe_concurrency = probe_concurrency.min(config.global_concurrency).max(1);
|
||||||
|
let probe_results = stream::iter(keys.into_iter().map(|key| {
|
||||||
|
let key_id = key.id.clone();
|
||||||
|
let admin_state = &admin_state;
|
||||||
|
let provider = &provider;
|
||||||
|
let endpoint = &endpoint;
|
||||||
|
let provider_type = provider_type.as_str();
|
||||||
|
async move {
|
||||||
|
let result = refresh_provider_probe_keys(
|
||||||
|
admin_state,
|
||||||
|
provider,
|
||||||
|
endpoint,
|
||||||
|
provider_type,
|
||||||
|
vec![key],
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
(key_id, result)
|
||||||
|
}
|
||||||
|
}))
|
||||||
|
.buffer_unordered(probe_concurrency)
|
||||||
|
.collect::<Vec<_>>()
|
||||||
|
.await;
|
||||||
|
|
||||||
|
let mut probe_success = 0usize;
|
||||||
|
let mut probe_failed = 0usize;
|
||||||
|
for (key_id, result) in probe_results {
|
||||||
|
match result {
|
||||||
|
Ok(payload) => {
|
||||||
|
update_summary_from_payload(&mut summary, 1, payload.as_ref());
|
||||||
|
probe_success += payload
|
||||||
|
.as_ref()
|
||||||
|
.and_then(|value| value.get("success"))
|
||||||
|
.and_then(Value::as_u64)
|
||||||
|
.unwrap_or(0) as usize;
|
||||||
|
probe_failed += payload
|
||||||
|
.as_ref()
|
||||||
|
.and_then(|value| value.get("failed"))
|
||||||
|
.and_then(Value::as_u64)
|
||||||
|
.unwrap_or(0) as usize;
|
||||||
|
record_score_probe_results_from_payload(
|
||||||
|
state,
|
||||||
|
&provider.id,
|
||||||
|
std::slice::from_ref(&key_id),
|
||||||
|
payload.as_ref(),
|
||||||
|
now_ts,
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
}
|
||||||
|
Err(err) => {
|
||||||
|
summary.failed += 1;
|
||||||
|
probe_failed += 1;
|
||||||
|
record_score_probe_result_for_key(
|
||||||
|
state,
|
||||||
|
&provider.id,
|
||||||
|
&key_id,
|
||||||
|
now_ts,
|
||||||
|
false,
|
||||||
|
Some(PoolMemberHardState::Cooldown),
|
||||||
|
serde_json::json!({
|
||||||
|
"last_probe": {
|
||||||
|
"source": "pool_quota_probe",
|
||||||
|
"status": "worker_error",
|
||||||
|
"message": format!("{err:?}")
|
||||||
|
}
|
||||||
|
}),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
warn!(
|
||||||
|
provider_id = %provider_short_id,
|
||||||
|
provider_type,
|
||||||
|
key_id,
|
||||||
|
error = ?err,
|
||||||
|
"gateway pool quota probe failed"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
info!(
|
||||||
|
provider_id = %provider_short_id,
|
||||||
|
provider_type,
|
||||||
|
selected = selected_count,
|
||||||
|
success = probe_success,
|
||||||
|
failed = probe_failed,
|
||||||
|
concurrency = probe_concurrency,
|
||||||
|
"gateway pool quota probe completed"
|
||||||
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
Ok(summary)
|
Ok(summary)
|
||||||
|
|||||||
@@ -0,0 +1,543 @@
|
|||||||
|
use std::collections::BTreeMap;
|
||||||
|
use std::time::{Duration, SystemTime, UNIX_EPOCH};
|
||||||
|
|
||||||
|
use aether_data_contracts::repository::global_models::AdminProviderModelListQuery;
|
||||||
|
use aether_data_contracts::repository::pool_scores::GetPoolMemberScoresByIdsQuery;
|
||||||
|
use aether_data_contracts::repository::provider_catalog::{
|
||||||
|
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
|
||||||
|
};
|
||||||
|
use tracing::{debug, info, warn};
|
||||||
|
|
||||||
|
use crate::admin_api::admin_provider_pool_config;
|
||||||
|
use crate::ai_serving::build_provider_key_pool_score_upsert;
|
||||||
|
use crate::handlers::shared::provider_pool::AdminProviderPoolConfig;
|
||||||
|
use crate::{AppState, GatewayError};
|
||||||
|
|
||||||
|
const POOL_SCORE_REBUILD_DEFAULT_INTERVAL_SECONDS: u64 = 300;
|
||||||
|
const POOL_SCORE_REBUILD_MIN_INTERVAL_SECONDS: u64 = 30;
|
||||||
|
const POOL_SCORE_REBUILD_DEFAULT_MAX_UPSERTS_PER_TICK: usize = 20_000;
|
||||||
|
const POOL_SCORE_REBUILD_PROVIDER_CURSOR_KEY: &str = "ap:pool_score_rebuild:provider_cursor";
|
||||||
|
const POOL_SCORE_REBUILD_PROVIDER_OFFSET_PREFIX: &str = "ap:pool_score_rebuild:provider_offset";
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||||
|
pub(crate) struct PoolScoreRebuildRunSummary {
|
||||||
|
pub(crate) providers_checked: usize,
|
||||||
|
pub(crate) providers_scored: usize,
|
||||||
|
pub(crate) keys_seen: usize,
|
||||||
|
pub(crate) scores_upserted: usize,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl PoolScoreRebuildRunSummary {
|
||||||
|
const fn empty() -> Self {
|
||||||
|
Self {
|
||||||
|
providers_checked: 0,
|
||||||
|
providers_scored: 0,
|
||||||
|
keys_seen: 0,
|
||||||
|
scores_upserted: 0,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Copy)]
|
||||||
|
pub(crate) struct PoolScoreRebuildWorkerConfig {
|
||||||
|
pub(crate) interval: Duration,
|
||||||
|
pub(crate) max_upserts_per_tick: usize,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl PoolScoreRebuildWorkerConfig {
|
||||||
|
fn from_env() -> Self {
|
||||||
|
let interval_seconds = env_u64(
|
||||||
|
"POOL_SCORE_REBUILD_INTERVAL_SECONDS",
|
||||||
|
POOL_SCORE_REBUILD_DEFAULT_INTERVAL_SECONDS,
|
||||||
|
)
|
||||||
|
.max(POOL_SCORE_REBUILD_MIN_INTERVAL_SECONDS);
|
||||||
|
let max_upserts_per_tick = env_usize(
|
||||||
|
"POOL_SCORE_REBUILD_MAX_UPSERTS_PER_TICK",
|
||||||
|
POOL_SCORE_REBUILD_DEFAULT_MAX_UPSERTS_PER_TICK,
|
||||||
|
)
|
||||||
|
.max(1);
|
||||||
|
Self {
|
||||||
|
interval: Duration::from_secs(interval_seconds),
|
||||||
|
max_upserts_per_tick,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn env_u64(name: &str, default_value: u64) -> u64 {
|
||||||
|
std::env::var(name)
|
||||||
|
.ok()
|
||||||
|
.and_then(|value| value.trim().parse::<u64>().ok())
|
||||||
|
.unwrap_or(default_value)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn env_usize(name: &str, default_value: usize) -> usize {
|
||||||
|
std::env::var(name)
|
||||||
|
.ok()
|
||||||
|
.and_then(|value| value.trim().parse::<usize>().ok())
|
||||||
|
.unwrap_or(default_value)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn now_unix_secs() -> u64 {
|
||||||
|
SystemTime::now()
|
||||||
|
.duration_since(UNIX_EPOCH)
|
||||||
|
.unwrap_or_default()
|
||||||
|
.as_secs()
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn load_runtime_usize(state: &AppState, key: &str) -> usize {
|
||||||
|
state
|
||||||
|
.runtime_state
|
||||||
|
.kv_get(key)
|
||||||
|
.await
|
||||||
|
.ok()
|
||||||
|
.flatten()
|
||||||
|
.and_then(|value| value.trim().parse::<usize>().ok())
|
||||||
|
.unwrap_or(0)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn store_runtime_usize(state: &AppState, key: &str, value: usize) {
|
||||||
|
if let Err(err) = state
|
||||||
|
.runtime_state
|
||||||
|
.kv_set(key, value.to_string(), None)
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
debug!(
|
||||||
|
key,
|
||||||
|
error = ?err,
|
||||||
|
"gateway pool score rebuild: failed to store cursor"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn provider_offset_cursor_key(provider_id: &str) -> String {
|
||||||
|
format!("{POOL_SCORE_REBUILD_PROVIDER_OFFSET_PREFIX}:{provider_id}")
|
||||||
|
}
|
||||||
|
|
||||||
|
fn score_combo_indices(
|
||||||
|
flat_index: usize,
|
||||||
|
model_count: usize,
|
||||||
|
key_count: usize,
|
||||||
|
) -> (usize, usize, usize) {
|
||||||
|
let endpoint_stride = model_count.saturating_mul(key_count).max(1);
|
||||||
|
let endpoint_index = flat_index / endpoint_stride;
|
||||||
|
let remainder = flat_index % endpoint_stride;
|
||||||
|
let model_index = remainder / key_count.max(1);
|
||||||
|
let key_index = remainder % key_count.max(1);
|
||||||
|
(endpoint_index, model_index, key_index)
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
struct ProviderScoreBuildItem {
|
||||||
|
endpoint_index: usize,
|
||||||
|
model_index: usize,
|
||||||
|
key_index: usize,
|
||||||
|
score_id: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(crate) async fn ensure_provider_key_pool_scores_for_keys(
|
||||||
|
state: &AppState,
|
||||||
|
provider: &StoredProviderCatalogProvider,
|
||||||
|
pool_config: &AdminProviderPoolConfig,
|
||||||
|
endpoints: &[StoredProviderCatalogEndpoint],
|
||||||
|
keys: &[StoredProviderCatalogKey],
|
||||||
|
now_unix_secs: u64,
|
||||||
|
max_upserts: usize,
|
||||||
|
) -> Result<usize, GatewayError> {
|
||||||
|
if max_upserts == 0
|
||||||
|
|| keys.is_empty()
|
||||||
|
|| !state.data.has_pool_score_reader()
|
||||||
|
|| !state.data.has_pool_score_writer()
|
||||||
|
{
|
||||||
|
return Ok(0);
|
||||||
|
}
|
||||||
|
|
||||||
|
let endpoints = endpoints
|
||||||
|
.iter()
|
||||||
|
.filter(|endpoint| endpoint.is_active && !endpoint.api_format.trim().is_empty())
|
||||||
|
.collect::<Vec<_>>();
|
||||||
|
let keys = keys
|
||||||
|
.iter()
|
||||||
|
.filter(|key| key.is_active && key.provider_id == provider.id)
|
||||||
|
.collect::<Vec<_>>();
|
||||||
|
if endpoints.is_empty() || keys.is_empty() {
|
||||||
|
return Ok(0);
|
||||||
|
}
|
||||||
|
|
||||||
|
let models = state
|
||||||
|
.list_admin_provider_models(&AdminProviderModelListQuery {
|
||||||
|
provider_id: provider.id.clone(),
|
||||||
|
is_active: Some(true),
|
||||||
|
offset: 0,
|
||||||
|
limit: 10_000,
|
||||||
|
})
|
||||||
|
.await?
|
||||||
|
.into_iter()
|
||||||
|
.filter(|model| model.is_available)
|
||||||
|
.collect::<Vec<_>>();
|
||||||
|
if models.is_empty() {
|
||||||
|
return Ok(0);
|
||||||
|
}
|
||||||
|
|
||||||
|
let max_items = endpoints
|
||||||
|
.len()
|
||||||
|
.saturating_mul(models.len())
|
||||||
|
.saturating_mul(keys.len())
|
||||||
|
.min(max_upserts);
|
||||||
|
let mut build_items = Vec::with_capacity(max_items);
|
||||||
|
'outer: for (endpoint_index, endpoint) in endpoints.iter().enumerate() {
|
||||||
|
for (model_index, model) in models.iter().enumerate() {
|
||||||
|
for (key_index, key) in keys.iter().enumerate() {
|
||||||
|
let draft = build_provider_key_pool_score_upsert(
|
||||||
|
key,
|
||||||
|
provider.provider_type.as_str(),
|
||||||
|
endpoint.api_format.trim(),
|
||||||
|
Some(model.id.as_str()),
|
||||||
|
None,
|
||||||
|
now_unix_secs,
|
||||||
|
pool_config.score_rules,
|
||||||
|
);
|
||||||
|
build_items.push(ProviderScoreBuildItem {
|
||||||
|
endpoint_index,
|
||||||
|
model_index,
|
||||||
|
key_index,
|
||||||
|
score_id: draft.id,
|
||||||
|
});
|
||||||
|
if build_items.len() >= max_upserts {
|
||||||
|
break 'outer;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if build_items.is_empty() {
|
||||||
|
return Ok(0);
|
||||||
|
}
|
||||||
|
|
||||||
|
let existing_score_ids = state
|
||||||
|
.data
|
||||||
|
.get_pool_member_scores_by_ids(&GetPoolMemberScoresByIdsQuery {
|
||||||
|
ids: build_items
|
||||||
|
.iter()
|
||||||
|
.map(|item| item.score_id.clone())
|
||||||
|
.collect(),
|
||||||
|
})
|
||||||
|
.await
|
||||||
|
.unwrap_or_else(|err| {
|
||||||
|
debug!(
|
||||||
|
provider_id = %provider.id,
|
||||||
|
error = ?err,
|
||||||
|
"gateway pool score ensure: failed to read existing scores by id"
|
||||||
|
);
|
||||||
|
Vec::new()
|
||||||
|
})
|
||||||
|
.into_iter()
|
||||||
|
.map(|score| score.id)
|
||||||
|
.collect::<std::collections::BTreeSet<_>>();
|
||||||
|
|
||||||
|
let mut upserted = 0usize;
|
||||||
|
for item in &build_items {
|
||||||
|
if existing_score_ids.contains(&item.score_id) {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
let endpoint = endpoints[item.endpoint_index];
|
||||||
|
let model = &models[item.model_index];
|
||||||
|
let key = keys[item.key_index];
|
||||||
|
let upsert = build_provider_key_pool_score_upsert(
|
||||||
|
key,
|
||||||
|
provider.provider_type.as_str(),
|
||||||
|
endpoint.api_format.trim(),
|
||||||
|
Some(model.id.as_str()),
|
||||||
|
None,
|
||||||
|
now_unix_secs,
|
||||||
|
pool_config.score_rules,
|
||||||
|
);
|
||||||
|
if state
|
||||||
|
.data
|
||||||
|
.upsert_pool_member_score(upsert)
|
||||||
|
.await
|
||||||
|
.map_err(|err| GatewayError::Internal(format!("{err:?}")))?
|
||||||
|
.is_some()
|
||||||
|
{
|
||||||
|
upserted = upserted.saturating_add(1);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(upserted)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(crate) async fn perform_pool_score_rebuild_once_with_config(
|
||||||
|
state: &AppState,
|
||||||
|
config: PoolScoreRebuildWorkerConfig,
|
||||||
|
) -> Result<PoolScoreRebuildRunSummary, GatewayError> {
|
||||||
|
if !state.has_provider_catalog_data_reader()
|
||||||
|
|| !state.data.has_pool_score_reader()
|
||||||
|
|| !state.data.has_pool_score_writer()
|
||||||
|
{
|
||||||
|
return Ok(PoolScoreRebuildRunSummary::empty());
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut providers = state
|
||||||
|
.list_provider_catalog_providers(true)
|
||||||
|
.await?
|
||||||
|
.into_iter()
|
||||||
|
.filter_map(|provider| {
|
||||||
|
admin_provider_pool_config(&provider).map(|config| (provider, config))
|
||||||
|
})
|
||||||
|
.collect::<Vec<_>>();
|
||||||
|
providers.sort_by(|left, right| left.0.id.cmp(&right.0.id));
|
||||||
|
if providers.is_empty() {
|
||||||
|
return Ok(PoolScoreRebuildRunSummary::empty());
|
||||||
|
}
|
||||||
|
|
||||||
|
let provider_ids = providers
|
||||||
|
.iter()
|
||||||
|
.map(|(provider, _)| provider.id.clone())
|
||||||
|
.collect::<Vec<_>>();
|
||||||
|
let mut endpoints_by_provider = BTreeMap::new();
|
||||||
|
for endpoint in state
|
||||||
|
.list_provider_catalog_endpoints_by_provider_ids(&provider_ids)
|
||||||
|
.await?
|
||||||
|
.into_iter()
|
||||||
|
.filter(|endpoint| endpoint.is_active)
|
||||||
|
{
|
||||||
|
endpoints_by_provider
|
||||||
|
.entry(endpoint.provider_id.clone())
|
||||||
|
.or_insert_with(Vec::new)
|
||||||
|
.push(endpoint);
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut keys_by_provider = BTreeMap::new();
|
||||||
|
for key in state
|
||||||
|
.list_provider_catalog_keys_by_provider_ids(&provider_ids)
|
||||||
|
.await?
|
||||||
|
{
|
||||||
|
keys_by_provider
|
||||||
|
.entry(key.provider_id.clone())
|
||||||
|
.or_insert_with(Vec::new)
|
||||||
|
.push(key);
|
||||||
|
}
|
||||||
|
|
||||||
|
let now = now_unix_secs();
|
||||||
|
let mut summary = PoolScoreRebuildRunSummary {
|
||||||
|
providers_checked: providers.len(),
|
||||||
|
..PoolScoreRebuildRunSummary::empty()
|
||||||
|
};
|
||||||
|
|
||||||
|
let start_provider_index =
|
||||||
|
load_runtime_usize(state, POOL_SCORE_REBUILD_PROVIDER_CURSOR_KEY).await % providers.len();
|
||||||
|
let mut last_provider_index = None;
|
||||||
|
for provider_index in
|
||||||
|
(0..providers.len()).map(|offset| (start_provider_index + offset) % providers.len())
|
||||||
|
{
|
||||||
|
if summary.scores_upserted >= config.max_upserts_per_tick {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
last_provider_index = Some(provider_index);
|
||||||
|
let (provider, pool_config) = providers[provider_index].clone();
|
||||||
|
let endpoints = endpoints_by_provider
|
||||||
|
.remove(&provider.id)
|
||||||
|
.unwrap_or_default();
|
||||||
|
let keys = keys_by_provider.remove(&provider.id).unwrap_or_default();
|
||||||
|
if endpoints.is_empty() || keys.is_empty() {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
let models = state
|
||||||
|
.list_admin_provider_models(&AdminProviderModelListQuery {
|
||||||
|
provider_id: provider.id.clone(),
|
||||||
|
is_active: Some(true),
|
||||||
|
offset: 0,
|
||||||
|
limit: 10_000,
|
||||||
|
})
|
||||||
|
.await?
|
||||||
|
.into_iter()
|
||||||
|
.filter(|model| model.is_available)
|
||||||
|
.collect::<Vec<_>>();
|
||||||
|
if models.is_empty() {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
let total_combinations = endpoints
|
||||||
|
.len()
|
||||||
|
.saturating_mul(models.len())
|
||||||
|
.saturating_mul(keys.len());
|
||||||
|
if total_combinations == 0 {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
let provider_cursor_key = provider_offset_cursor_key(&provider.id);
|
||||||
|
let provider_cursor =
|
||||||
|
load_runtime_usize(state, &provider_cursor_key).await % total_combinations.max(1);
|
||||||
|
let remaining_budget = config
|
||||||
|
.max_upserts_per_tick
|
||||||
|
.saturating_sub(summary.scores_upserted);
|
||||||
|
let provider_budget = remaining_budget.min(total_combinations);
|
||||||
|
let mut build_items = Vec::with_capacity(provider_budget);
|
||||||
|
for offset in 0..provider_budget {
|
||||||
|
let flat_index = (provider_cursor + offset) % total_combinations;
|
||||||
|
let (endpoint_index, model_index, key_index) =
|
||||||
|
score_combo_indices(flat_index, models.len(), keys.len());
|
||||||
|
let endpoint = &endpoints[endpoint_index];
|
||||||
|
let api_format = endpoint.api_format.trim();
|
||||||
|
if api_format.is_empty() {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
let model = &models[model_index];
|
||||||
|
let key = &keys[key_index];
|
||||||
|
let draft = build_provider_key_pool_score_upsert(
|
||||||
|
key,
|
||||||
|
provider.provider_type.as_str(),
|
||||||
|
api_format,
|
||||||
|
Some(model.id.as_str()),
|
||||||
|
None,
|
||||||
|
now,
|
||||||
|
pool_config.score_rules,
|
||||||
|
);
|
||||||
|
build_items.push(ProviderScoreBuildItem {
|
||||||
|
endpoint_index,
|
||||||
|
model_index,
|
||||||
|
key_index,
|
||||||
|
score_id: draft.id,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
if build_items.is_empty() {
|
||||||
|
store_runtime_usize(
|
||||||
|
state,
|
||||||
|
&provider_cursor_key,
|
||||||
|
(provider_cursor + provider_budget) % total_combinations,
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
let existing_scores = state
|
||||||
|
.data
|
||||||
|
.get_pool_member_scores_by_ids(&GetPoolMemberScoresByIdsQuery {
|
||||||
|
ids: build_items
|
||||||
|
.iter()
|
||||||
|
.map(|item| item.score_id.clone())
|
||||||
|
.collect(),
|
||||||
|
})
|
||||||
|
.await
|
||||||
|
.unwrap_or_else(|err| {
|
||||||
|
debug!(
|
||||||
|
provider_id = %provider.id,
|
||||||
|
error = ?err,
|
||||||
|
"gateway pool score rebuild: failed to read existing scores by id"
|
||||||
|
);
|
||||||
|
Vec::new()
|
||||||
|
})
|
||||||
|
.into_iter()
|
||||||
|
.map(|score| (score.id.clone(), score))
|
||||||
|
.collect::<BTreeMap<_, _>>();
|
||||||
|
let mut provider_upserts = 0usize;
|
||||||
|
summary.keys_seen = summary.keys_seen.saturating_add(keys.len());
|
||||||
|
for item in &build_items {
|
||||||
|
if summary.scores_upserted >= config.max_upserts_per_tick {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
let endpoint = &endpoints[item.endpoint_index];
|
||||||
|
let model = &models[item.model_index];
|
||||||
|
let key = &keys[item.key_index];
|
||||||
|
let existing = existing_scores.get(&item.score_id);
|
||||||
|
let upsert = build_provider_key_pool_score_upsert(
|
||||||
|
key,
|
||||||
|
provider.provider_type.as_str(),
|
||||||
|
endpoint.api_format.trim(),
|
||||||
|
Some(model.id.as_str()),
|
||||||
|
existing,
|
||||||
|
now,
|
||||||
|
pool_config.score_rules,
|
||||||
|
);
|
||||||
|
if state
|
||||||
|
.data
|
||||||
|
.upsert_pool_member_score(upsert)
|
||||||
|
.await
|
||||||
|
.map_err(|err| GatewayError::Internal(format!("{err:?}")))?
|
||||||
|
.is_some()
|
||||||
|
{
|
||||||
|
summary.scores_upserted = summary.scores_upserted.saturating_add(1);
|
||||||
|
provider_upserts = provider_upserts.saturating_add(1);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
store_runtime_usize(
|
||||||
|
state,
|
||||||
|
&provider_cursor_key,
|
||||||
|
(provider_cursor + provider_budget) % total_combinations,
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
if provider_upserts > 0 {
|
||||||
|
summary.providers_scored = summary.providers_scored.saturating_add(1);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if let Some(last_provider_index) = last_provider_index {
|
||||||
|
store_runtime_usize(
|
||||||
|
state,
|
||||||
|
POOL_SCORE_REBUILD_PROVIDER_CURSOR_KEY,
|
||||||
|
(last_provider_index + 1) % providers.len(),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(summary)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(crate) async fn perform_pool_score_rebuild_once(
|
||||||
|
state: &AppState,
|
||||||
|
) -> Result<PoolScoreRebuildRunSummary, GatewayError> {
|
||||||
|
perform_pool_score_rebuild_once_with_config(state, PoolScoreRebuildWorkerConfig::from_env())
|
||||||
|
.await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(crate) fn spawn_pool_score_rebuild_worker(
|
||||||
|
state: AppState,
|
||||||
|
) -> Option<tokio::task::JoinHandle<()>> {
|
||||||
|
if !state.has_provider_catalog_data_reader()
|
||||||
|
|| !state.data.has_pool_score_reader()
|
||||||
|
|| !state.data.has_pool_score_writer()
|
||||||
|
{
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
|
||||||
|
let config = PoolScoreRebuildWorkerConfig::from_env();
|
||||||
|
Some(tokio::spawn(async move {
|
||||||
|
if let Err(err) = perform_pool_score_rebuild_once_with_config(&state, config).await {
|
||||||
|
warn!(
|
||||||
|
error = ?err,
|
||||||
|
"gateway pool score rebuild initial tick failed"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
let mut interval = tokio::time::interval(config.interval);
|
||||||
|
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
|
||||||
|
loop {
|
||||||
|
interval.tick().await;
|
||||||
|
match perform_pool_score_rebuild_once_with_config(&state, config).await {
|
||||||
|
Ok(summary) if summary.scores_upserted > 0 => {
|
||||||
|
info!(
|
||||||
|
providers_checked = summary.providers_checked,
|
||||||
|
providers_scored = summary.providers_scored,
|
||||||
|
keys_seen = summary.keys_seen,
|
||||||
|
scores_upserted = summary.scores_upserted,
|
||||||
|
"gateway pool score rebuild completed"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
Ok(_) => {}
|
||||||
|
Err(err) => {
|
||||||
|
warn!(
|
||||||
|
error = ?err,
|
||||||
|
"gateway pool score rebuild worker tick failed"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}))
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::score_combo_indices;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn score_combo_indices_walks_endpoint_model_key_order() {
|
||||||
|
assert_eq!(score_combo_indices(0, 2, 3), (0, 0, 0));
|
||||||
|
assert_eq!(score_combo_indices(2, 2, 3), (0, 0, 2));
|
||||||
|
assert_eq!(score_combo_indices(3, 2, 3), (0, 1, 0));
|
||||||
|
assert_eq!(score_combo_indices(6, 2, 3), (1, 0, 0));
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -2,6 +2,9 @@ use std::collections::BTreeMap;
|
|||||||
|
|
||||||
use aether_admin::provider::quota as admin_provider_quota_pure;
|
use aether_admin::provider::quota as admin_provider_quota_pure;
|
||||||
use aether_contracts::{ExecutionPlan, ExecutionTelemetry};
|
use aether_contracts::{ExecutionPlan, ExecutionTelemetry};
|
||||||
|
use aether_data_contracts::repository::pool_scores::{
|
||||||
|
PoolMemberHardState, PoolMemberIdentity, PoolMemberScheduleFeedback,
|
||||||
|
};
|
||||||
use aether_scheduler_core::{
|
use aether_scheduler_core::{
|
||||||
build_scheduler_affinity_cache_key_for_api_key_id_with_client_session,
|
build_scheduler_affinity_cache_key_for_api_key_id_with_client_session,
|
||||||
count_recent_rpm_requests_for_provider_key, ClientSessionAffinity, SchedulerAffinityTarget,
|
count_recent_rpm_requests_for_provider_key, ClientSessionAffinity, SchedulerAffinityTarget,
|
||||||
@@ -381,6 +384,19 @@ async fn record_sync_pool_success_effect(
|
|||||||
resolve_ttfb_ms(payload.telemetry.as_ref()),
|
resolve_ttfb_ms(payload.telemetry.as_ref()),
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
|
record_pool_score_schedule_feedback(
|
||||||
|
state,
|
||||||
|
context,
|
||||||
|
Some(true),
|
||||||
|
Some(PoolMemberHardState::Available),
|
||||||
|
Some(50),
|
||||||
|
serde_json::json!({
|
||||||
|
"last_request_feedback": {
|
||||||
|
"source": "sync_success"
|
||||||
|
}
|
||||||
|
}),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn record_adaptive_rate_limit_effect(
|
async fn record_adaptive_rate_limit_effect(
|
||||||
@@ -600,6 +616,19 @@ async fn record_stream_pool_success_effect(
|
|||||||
resolve_ttfb_ms(payload.telemetry.as_ref()),
|
resolve_ttfb_ms(payload.telemetry.as_ref()),
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
|
record_pool_score_schedule_feedback(
|
||||||
|
state,
|
||||||
|
context,
|
||||||
|
Some(true),
|
||||||
|
Some(PoolMemberHardState::Available),
|
||||||
|
Some(50),
|
||||||
|
serde_json::json!({
|
||||||
|
"last_request_feedback": {
|
||||||
|
"source": "stream_success"
|
||||||
|
}
|
||||||
|
}),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn record_pool_error_effect(
|
async fn record_pool_error_effect(
|
||||||
@@ -636,6 +665,21 @@ async fn record_pool_error_effect(
|
|||||||
Some(effect.headers),
|
Some(effect.headers),
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
|
record_pool_score_schedule_feedback(
|
||||||
|
state,
|
||||||
|
context,
|
||||||
|
Some(false),
|
||||||
|
pool_score_hard_state_for_status(effect.status_code, effect.error_body),
|
||||||
|
Some(pool_score_delta_for_status(effect.status_code)),
|
||||||
|
serde_json::json!({
|
||||||
|
"last_request_feedback": {
|
||||||
|
"source": "pool_error",
|
||||||
|
"status_code": effect.status_code,
|
||||||
|
"classification": format!("{:?}", effect.classification)
|
||||||
|
}
|
||||||
|
}),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn open_pool_key_circuit_breaker(
|
async fn open_pool_key_circuit_breaker(
|
||||||
@@ -730,6 +774,21 @@ async fn record_oauth_invalidation_effect(
|
|||||||
plan.provider_id, plan.endpoint_id, plan.key_id, err
|
plan.provider_id, plan.endpoint_id, plan.key_id, err
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
record_pool_score_schedule_feedback(
|
||||||
|
state,
|
||||||
|
context,
|
||||||
|
Some(false),
|
||||||
|
Some(PoolMemberHardState::AuthInvalid),
|
||||||
|
Some(-2_000),
|
||||||
|
serde_json::json!({
|
||||||
|
"last_request_feedback": {
|
||||||
|
"source": "oauth_invalidation",
|
||||||
|
"status_code": effect.status_code,
|
||||||
|
"reason": invalid_reason
|
||||||
|
}
|
||||||
|
}),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
}
|
}
|
||||||
|
|
||||||
fn resolve_local_oauth_invalid_reason(
|
fn resolve_local_oauth_invalid_reason(
|
||||||
@@ -792,6 +851,92 @@ async fn record_pool_stream_timeout_effect(
|
|||||||
&pool_context.pool_config,
|
&pool_context.pool_config,
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
|
record_pool_score_schedule_feedback(
|
||||||
|
state,
|
||||||
|
context,
|
||||||
|
Some(false),
|
||||||
|
Some(PoolMemberHardState::Cooldown),
|
||||||
|
Some(-250),
|
||||||
|
serde_json::json!({
|
||||||
|
"last_request_feedback": {
|
||||||
|
"source": "stream_timeout"
|
||||||
|
}
|
||||||
|
}),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn record_pool_score_schedule_feedback(
|
||||||
|
state: &AppState,
|
||||||
|
context: LocalExecutionEffectContext<'_>,
|
||||||
|
succeeded: Option<bool>,
|
||||||
|
hard_state: Option<PoolMemberHardState>,
|
||||||
|
score_delta: Option<i32>,
|
||||||
|
score_reason_patch: Value,
|
||||||
|
) {
|
||||||
|
if context.plan.provider_id.trim().is_empty() || context.plan.key_id.trim().is_empty() {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
let feedback = PoolMemberScheduleFeedback {
|
||||||
|
identity: PoolMemberIdentity::provider_api_key(
|
||||||
|
context.plan.provider_id.clone(),
|
||||||
|
context.plan.key_id.clone(),
|
||||||
|
),
|
||||||
|
scope: None,
|
||||||
|
scheduled_at: current_unix_secs(),
|
||||||
|
succeeded,
|
||||||
|
hard_state,
|
||||||
|
score_delta,
|
||||||
|
score_reason_patch: Some(score_reason_patch),
|
||||||
|
};
|
||||||
|
if let Err(err) = state
|
||||||
|
.data
|
||||||
|
.record_pool_member_schedule_feedback(feedback)
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
warn!(
|
||||||
|
provider_id = %context.plan.provider_id,
|
||||||
|
key_id = %context.plan.key_id,
|
||||||
|
error = ?err,
|
||||||
|
"gateway orchestration effects: failed to record pool score schedule feedback"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn pool_score_hard_state_for_status(
|
||||||
|
status_code: u16,
|
||||||
|
error_body: Option<&str>,
|
||||||
|
) -> Option<PoolMemberHardState> {
|
||||||
|
match status_code {
|
||||||
|
401 | 403 => Some(PoolMemberHardState::AuthInvalid),
|
||||||
|
402 => Some(PoolMemberHardState::QuotaExhausted),
|
||||||
|
429 | 500..=599 => Some(PoolMemberHardState::Cooldown),
|
||||||
|
_ => {
|
||||||
|
let body = error_body.unwrap_or_default().to_ascii_lowercase();
|
||||||
|
if body.contains("quota") && body.contains("exceed") {
|
||||||
|
Some(PoolMemberHardState::QuotaExhausted)
|
||||||
|
} else if body.contains("invalid") && body.contains("token") {
|
||||||
|
Some(PoolMemberHardState::AuthInvalid)
|
||||||
|
} else if body.contains("banned")
|
||||||
|
|| body.contains("suspended")
|
||||||
|
|| body.contains("blocked")
|
||||||
|
{
|
||||||
|
Some(PoolMemberHardState::Banned)
|
||||||
|
} else {
|
||||||
|
None
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn pool_score_delta_for_status(status_code: u16) -> i32 {
|
||||||
|
match status_code {
|
||||||
|
401 | 403 => -2_000,
|
||||||
|
402 => -1_000,
|
||||||
|
429 => -500,
|
||||||
|
500..=599 => -300,
|
||||||
|
_ => -100,
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
|
|||||||
@@ -1,7 +1,8 @@
|
|||||||
use super::{AppState, GatewayError, LocalMutationOutcome, LocalProviderDeleteTaskState};
|
use super::{AppState, GatewayError, LocalMutationOutcome, LocalProviderDeleteTaskState};
|
||||||
use crate::handlers::shared::sync_provider_key_oauth_status_snapshot;
|
use crate::handlers::shared::sync_provider_key_oauth_status_snapshot;
|
||||||
use aether_data_contracts::repository::{candidates, global_models, provider_catalog};
|
use aether_data_contracts::repository::{candidates, global_models, pool_scores, provider_catalog};
|
||||||
use std::time::{SystemTime, UNIX_EPOCH};
|
use std::time::{SystemTime, UNIX_EPOCH};
|
||||||
|
use tracing::warn;
|
||||||
|
|
||||||
impl AppState {
|
impl AppState {
|
||||||
pub fn has_provider_catalog_data_reader(&self) -> bool {
|
pub fn has_provider_catalog_data_reader(&self) -> bool {
|
||||||
@@ -529,6 +530,25 @@ impl AppState {
|
|||||||
.cleanup_deleted_provider_catalog_refs(provider_id, endpoint_ids, key_ids)
|
.cleanup_deleted_provider_catalog_refs(provider_id, endpoint_ids, key_ids)
|
||||||
.await
|
.await
|
||||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||||
|
for key_id in key_ids {
|
||||||
|
if let Err(err) = self
|
||||||
|
.data
|
||||||
|
.delete_pool_member_scores_for_member(
|
||||||
|
&pool_scores::PoolMemberIdentity::provider_api_key(
|
||||||
|
provider_id.to_string(),
|
||||||
|
key_id.to_string(),
|
||||||
|
),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
warn!(
|
||||||
|
provider_id,
|
||||||
|
key_id,
|
||||||
|
error = ?err,
|
||||||
|
"gateway provider catalog cleanup: failed to delete pool member scores"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
if !endpoint_ids.is_empty() || !key_ids.is_empty() {
|
if !endpoint_ids.is_empty() || !key_ids.is_empty() {
|
||||||
self.clear_provider_transport_snapshot_cache();
|
self.clear_provider_transport_snapshot_cache();
|
||||||
}
|
}
|
||||||
@@ -620,12 +640,38 @@ impl AppState {
|
|||||||
&self,
|
&self,
|
||||||
key_id: &str,
|
key_id: &str,
|
||||||
) -> Result<bool, GatewayError> {
|
) -> Result<bool, GatewayError> {
|
||||||
|
let existing_key = self
|
||||||
|
.data
|
||||||
|
.list_provider_catalog_keys_by_ids(&[key_id.to_string()])
|
||||||
|
.await
|
||||||
|
.map_err(|err| GatewayError::Internal(err.to_string()))?
|
||||||
|
.into_iter()
|
||||||
|
.next();
|
||||||
let deleted = self
|
let deleted = self
|
||||||
.data
|
.data
|
||||||
.delete_provider_catalog_key(key_id)
|
.delete_provider_catalog_key(key_id)
|
||||||
.await
|
.await
|
||||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||||
if deleted {
|
if deleted {
|
||||||
|
if let Some(key) = existing_key {
|
||||||
|
if let Err(err) = self
|
||||||
|
.data
|
||||||
|
.delete_pool_member_scores_for_member(
|
||||||
|
&pool_scores::PoolMemberIdentity::provider_api_key(
|
||||||
|
key.provider_id.clone(),
|
||||||
|
key.id.clone(),
|
||||||
|
),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
warn!(
|
||||||
|
provider_id = %key.provider_id,
|
||||||
|
key_id = %key.id,
|
||||||
|
error = ?err,
|
||||||
|
"gateway provider catalog key delete: failed to delete pool member scores"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
self.clear_provider_transport_snapshot_cache();
|
self.clear_provider_transport_snapshot_cache();
|
||||||
}
|
}
|
||||||
Ok(deleted)
|
Ok(deleted)
|
||||||
|
|||||||
@@ -47,6 +47,7 @@ use crate::maintenance::spawn_oauth_token_refresh_worker;
|
|||||||
use crate::maintenance::spawn_pending_cleanup_worker;
|
use crate::maintenance::spawn_pending_cleanup_worker;
|
||||||
use crate::maintenance::spawn_pool_monitor_worker;
|
use crate::maintenance::spawn_pool_monitor_worker;
|
||||||
use crate::maintenance::spawn_pool_quota_probe_worker;
|
use crate::maintenance::spawn_pool_quota_probe_worker;
|
||||||
|
use crate::maintenance::spawn_pool_score_rebuild_worker;
|
||||||
use crate::maintenance::spawn_provider_checkin_worker;
|
use crate::maintenance::spawn_provider_checkin_worker;
|
||||||
use crate::maintenance::spawn_proxy_node_metrics_cleanup_worker;
|
use crate::maintenance::spawn_proxy_node_metrics_cleanup_worker;
|
||||||
use crate::maintenance::spawn_proxy_node_stale_cleanup_worker;
|
use crate::maintenance::spawn_proxy_node_stale_cleanup_worker;
|
||||||
@@ -1113,6 +1114,10 @@ impl AppState {
|
|||||||
crate::task_runtime::TASK_KEY_POOL_QUOTA_PROBE,
|
crate::task_runtime::TASK_KEY_POOL_QUOTA_PROBE,
|
||||||
spawn_pool_quota_probe_worker(self.clone()),
|
spawn_pool_quota_probe_worker(self.clone()),
|
||||||
);
|
);
|
||||||
|
supervise_worker(
|
||||||
|
crate::task_runtime::TASK_KEY_POOL_SCORE_REBUILD,
|
||||||
|
spawn_pool_score_rebuild_worker(self.clone()),
|
||||||
|
);
|
||||||
supervise_worker(
|
supervise_worker(
|
||||||
crate::task_runtime::TASK_KEY_STATS_HOURLY_AGG,
|
crate::task_runtime::TASK_KEY_STATS_HOURLY_AGG,
|
||||||
spawn_stats_hourly_aggregation_worker(self.data.clone()),
|
spawn_stats_hourly_aggregation_worker(self.data.clone()),
|
||||||
|
|||||||
@@ -57,6 +57,16 @@ impl AppState {
|
|||||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub(crate) async fn list_pool_key_candidate_rows_for_group_key_ids(
|
||||||
|
&self,
|
||||||
|
query: &candidate_selection::StoredPoolKeyCandidateRowsByKeyIdsQuery,
|
||||||
|
) -> Result<Vec<candidate_selection::StoredMinimalCandidateSelectionRow>, GatewayError> {
|
||||||
|
self.data
|
||||||
|
.list_pool_key_candidate_rows_for_group_key_ids(query)
|
||||||
|
.await
|
||||||
|
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||||
|
}
|
||||||
|
|
||||||
pub(crate) async fn read_provider_quota_snapshot(
|
pub(crate) async fn read_provider_quota_snapshot(
|
||||||
&self,
|
&self,
|
||||||
provider_id: &str,
|
provider_id: &str,
|
||||||
|
|||||||
@@ -22,6 +22,7 @@ pub(crate) const TASK_KEY_VIDEO_TASK_POLLER: &str = "video.task.poller";
|
|||||||
pub(crate) const TASK_KEY_MODEL_FETCH_WORKER: &str = "model.fetch.worker";
|
pub(crate) const TASK_KEY_MODEL_FETCH_WORKER: &str = "model.fetch.worker";
|
||||||
pub(crate) const TASK_KEY_PROVIDER_QUOTA_RESET: &str = "provider.quota.reset.worker";
|
pub(crate) const TASK_KEY_PROVIDER_QUOTA_RESET: &str = "provider.quota.reset.worker";
|
||||||
pub(crate) const TASK_KEY_POOL_QUOTA_PROBE: &str = "pool.quota.probe.worker";
|
pub(crate) const TASK_KEY_POOL_QUOTA_PROBE: &str = "pool.quota.probe.worker";
|
||||||
|
pub(crate) const TASK_KEY_POOL_SCORE_REBUILD: &str = "pool.score.rebuild.worker";
|
||||||
pub(crate) const TASK_KEY_POOL_MONITOR: &str = "pool.monitor.worker";
|
pub(crate) const TASK_KEY_POOL_MONITOR: &str = "pool.monitor.worker";
|
||||||
pub(crate) const TASK_KEY_AUDIT_CLEANUP: &str = "maintenance.audit.cleanup";
|
pub(crate) const TASK_KEY_AUDIT_CLEANUP: &str = "maintenance.audit.cleanup";
|
||||||
pub(crate) const TASK_KEY_DB_MAINTENANCE: &str = "maintenance.database";
|
pub(crate) const TASK_KEY_DB_MAINTENANCE: &str = "maintenance.database";
|
||||||
@@ -101,6 +102,14 @@ const TASK_DEFINITIONS: &[TaskDefinition] = &[
|
|||||||
true,
|
true,
|
||||||
RETRY_ONCE,
|
RETRY_ONCE,
|
||||||
),
|
),
|
||||||
|
TaskDefinition::new(
|
||||||
|
TASK_KEY_POOL_SCORE_REBUILD,
|
||||||
|
TaskKind::Scheduled,
|
||||||
|
"interval",
|
||||||
|
true,
|
||||||
|
true,
|
||||||
|
RETRY_ONCE,
|
||||||
|
),
|
||||||
TaskDefinition::new(
|
TaskDefinition::new(
|
||||||
TASK_KEY_POOL_MONITOR,
|
TASK_KEY_POOL_MONITOR,
|
||||||
TaskKind::Scheduled,
|
TaskKind::Scheduled,
|
||||||
|
|||||||
@@ -1,8 +1,13 @@
|
|||||||
use std::sync::{Arc, Mutex};
|
use std::sync::{Arc, Mutex};
|
||||||
|
|
||||||
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
|
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
|
||||||
|
use aether_data::repository::pool_scores::InMemoryPoolMemberScoreRepository;
|
||||||
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
|
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
|
||||||
use aether_data::repository::usage::InMemoryUsageReadRepository;
|
use aether_data::repository::usage::InMemoryUsageReadRepository;
|
||||||
|
use aether_data_contracts::repository::pool_scores::{
|
||||||
|
PoolMemberHardState, PoolMemberProbeStatus, StoredPoolMemberScore, POOL_KIND_PROVIDER_KEY_POOL,
|
||||||
|
POOL_MEMBER_KIND_PROVIDER_API_KEY, POOL_SCORE_SCOPE_KIND_MODEL,
|
||||||
|
};
|
||||||
use aether_data_contracts::repository::provider_catalog::ProviderCatalogReadRepository;
|
use aether_data_contracts::repository::provider_catalog::ProviderCatalogReadRepository;
|
||||||
use aether_data_contracts::repository::usage::StoredRequestUsageAudit;
|
use aether_data_contracts::repository::usage::StoredRequestUsageAudit;
|
||||||
use axum::body::{to_bytes, Body, Bytes};
|
use axum::body::{to_bytes, Body, Bytes};
|
||||||
@@ -343,6 +348,94 @@ async fn gateway_handles_admin_pool_batch_import_locally_with_trusted_admin_prin
|
|||||||
upstream_handle.abort();
|
upstream_handle.abort();
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn gateway_handles_admin_pool_scores_locally_with_trusted_admin_principal() {
|
||||||
|
let provider = sample_provider("provider-openai", "openai", 10).with_transport_fields(
|
||||||
|
true,
|
||||||
|
false,
|
||||||
|
true,
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
Some(json!({
|
||||||
|
"pool_advanced": {
|
||||||
|
"enabled": true
|
||||||
|
}
|
||||||
|
})),
|
||||||
|
);
|
||||||
|
let mut key = sample_key("key-openai-a", "provider-openai", "openai:chat", "sk-a");
|
||||||
|
key.name = "score key".to_string();
|
||||||
|
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||||
|
vec![provider],
|
||||||
|
Vec::new(),
|
||||||
|
vec![key.clone()],
|
||||||
|
));
|
||||||
|
let pool_score_repository = Arc::new(InMemoryPoolMemberScoreRepository::seed(vec![
|
||||||
|
StoredPoolMemberScore {
|
||||||
|
id: "pms-provider-openai-key-openai-a-openai-chat-model-1".to_string(),
|
||||||
|
pool_kind: POOL_KIND_PROVIDER_KEY_POOL.to_string(),
|
||||||
|
pool_id: "provider-openai".to_string(),
|
||||||
|
member_kind: POOL_MEMBER_KIND_PROVIDER_API_KEY.to_string(),
|
||||||
|
member_id: "key-openai-a".to_string(),
|
||||||
|
capability: "openai:chat".to_string(),
|
||||||
|
scope_kind: POOL_SCORE_SCOPE_KIND_MODEL.to_string(),
|
||||||
|
scope_id: Some("model-1".to_string()),
|
||||||
|
score: 0.875,
|
||||||
|
hard_state: PoolMemberHardState::Available,
|
||||||
|
score_version: 1,
|
||||||
|
score_reason: json!({ "weights": { "manual_priority": 0.3 } }),
|
||||||
|
last_ranked_at: Some(1_700_000_000),
|
||||||
|
last_scheduled_at: Some(1_700_000_010),
|
||||||
|
last_success_at: Some(1_700_000_020),
|
||||||
|
last_failure_at: None,
|
||||||
|
failure_count: 0,
|
||||||
|
last_probe_attempt_at: Some(1_700_000_030),
|
||||||
|
last_probe_success_at: Some(1_700_000_040),
|
||||||
|
last_probe_failure_at: None,
|
||||||
|
probe_failure_count: 0,
|
||||||
|
probe_status: PoolMemberProbeStatus::Ok,
|
||||||
|
updated_at: 1_700_000_050,
|
||||||
|
},
|
||||||
|
]));
|
||||||
|
|
||||||
|
let gateway = build_router_with_state(
|
||||||
|
AppState::new()
|
||||||
|
.expect("gateway should build")
|
||||||
|
.with_data_state_for_tests(
|
||||||
|
GatewayDataState::with_provider_catalog_repository_for_tests(Arc::clone(
|
||||||
|
&provider_catalog_repository,
|
||||||
|
))
|
||||||
|
.with_pool_score_repository_for_tests(Arc::clone(&pool_score_repository)),
|
||||||
|
),
|
||||||
|
);
|
||||||
|
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||||
|
|
||||||
|
let response = reqwest::Client::new()
|
||||||
|
.get(format!(
|
||||||
|
"{gateway_url}/api/admin/pool/provider-openai/scores?api_format=openai:chat&model_id=model-1"
|
||||||
|
))
|
||||||
|
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||||
|
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||||
|
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||||
|
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.expect("request should succeed");
|
||||||
|
assert_eq!(response.status(), StatusCode::OK);
|
||||||
|
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||||
|
assert_eq!(
|
||||||
|
payload["items"].as_array().map(|items| items.len()),
|
||||||
|
Some(1)
|
||||||
|
);
|
||||||
|
assert_eq!(payload["items"][0]["member_id"], json!("key-openai-a"));
|
||||||
|
assert_eq!(payload["items"][0]["key"]["name"], json!("score key"));
|
||||||
|
assert_eq!(payload["items"][0]["probe_status"], json!("ok"));
|
||||||
|
|
||||||
|
gateway_handle.abort();
|
||||||
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn gateway_handles_admin_pool_trailing_slash_routes_locally_with_trusted_admin_principal() {
|
async fn gateway_handles_admin_pool_trailing_slash_routes_locally_with_trusted_admin_principal() {
|
||||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||||
|
|||||||
@@ -9,6 +9,7 @@ description = "AI serving application contracts and ports for Aether"
|
|||||||
[dependencies]
|
[dependencies]
|
||||||
aether-ai-formats.workspace = true
|
aether-ai-formats.workspace = true
|
||||||
aether-contracts.workspace = true
|
aether-contracts.workspace = true
|
||||||
|
aether-data-contracts.workspace = true
|
||||||
aether-scheduler-core.workspace = true
|
aether-scheduler-core.workspace = true
|
||||||
async-trait.workspace = true
|
async-trait.workspace = true
|
||||||
http.workspace = true
|
http.workspace = true
|
||||||
|
|||||||
@@ -16,6 +16,7 @@ pub mod execution_path;
|
|||||||
pub mod failure_diagnostic;
|
pub mod failure_diagnostic;
|
||||||
pub mod plan_payload;
|
pub mod plan_payload;
|
||||||
pub mod pool_scheduler;
|
pub mod pool_scheduler;
|
||||||
|
pub mod pool_scores;
|
||||||
pub mod ports;
|
pub mod ports;
|
||||||
pub mod ranking_metadata;
|
pub mod ranking_metadata;
|
||||||
pub mod report_context;
|
pub mod report_context;
|
||||||
@@ -99,6 +100,13 @@ pub use pool_scheduler::{
|
|||||||
AI_POOL_ACCOUNT_EXHAUSTED_SKIP_REASON, AI_POOL_COOLDOWN_SKIP_REASON,
|
AI_POOL_ACCOUNT_EXHAUSTED_SKIP_REASON, AI_POOL_COOLDOWN_SKIP_REASON,
|
||||||
AI_POOL_COST_LIMIT_REACHED_SKIP_REASON,
|
AI_POOL_COST_LIMIT_REACHED_SKIP_REASON,
|
||||||
};
|
};
|
||||||
|
pub use pool_scores::{
|
||||||
|
probe_freshness_score, probe_freshness_score_with_ttl, score_pool_member,
|
||||||
|
score_pool_member_with_rules, PoolMemberScoreInput, PoolMemberScoreOutput,
|
||||||
|
PoolMemberScoreRules, PoolMemberScoreWeights, POOL_SCORE_VERSION,
|
||||||
|
PROBE_FAILURE_COOLDOWN_THRESHOLD, PROBE_FAILURE_PENALTY, PROBE_FRESHNESS_TTL_SECONDS,
|
||||||
|
REQUEST_FAILURE_PENALTY, UNSCHEDULABLE_SCORE_CAP,
|
||||||
|
};
|
||||||
pub use ranking_metadata::append_ai_ranking_metadata_to_object;
|
pub use ranking_metadata::append_ai_ranking_metadata_to_object;
|
||||||
pub use report_context::{
|
pub use report_context::{
|
||||||
build_ai_execution_report_context, build_ai_report_context_original_request_echo,
|
build_ai_execution_report_context, build_ai_report_context_original_request_echo,
|
||||||
|
|||||||
540
crates/aether-ai-serving/src/pool_scores.rs
Normal file
540
crates/aether-ai-serving/src/pool_scores.rs
Normal file
@@ -0,0 +1,540 @@
|
|||||||
|
use aether_data_contracts::repository::pool_scores::{
|
||||||
|
PoolMemberHardState, PoolMemberIdentity, PoolMemberProbeStatus, PoolScoreScope,
|
||||||
|
};
|
||||||
|
use serde_json::{json, Value};
|
||||||
|
|
||||||
|
pub const POOL_SCORE_VERSION: u64 = 1;
|
||||||
|
pub const PROBE_FRESHNESS_TTL_SECONDS: u64 = 30 * 60;
|
||||||
|
pub const UNSCHEDULABLE_SCORE_CAP: f64 = 0.05;
|
||||||
|
pub const PROBE_FAILURE_PENALTY: f64 = 0.05;
|
||||||
|
pub const REQUEST_FAILURE_PENALTY: f64 = 0.005;
|
||||||
|
pub const PROBE_FAILURE_COOLDOWN_THRESHOLD: u64 = 3;
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Copy, PartialEq)]
|
||||||
|
pub struct PoolMemberScoreWeights {
|
||||||
|
pub manual_priority: f64,
|
||||||
|
pub health: f64,
|
||||||
|
pub probe_freshness: f64,
|
||||||
|
pub quota_remaining: f64,
|
||||||
|
pub latency: f64,
|
||||||
|
pub cost_lru: f64,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for PoolMemberScoreWeights {
|
||||||
|
fn default() -> Self {
|
||||||
|
Self {
|
||||||
|
manual_priority: 0.30,
|
||||||
|
health: 0.20,
|
||||||
|
probe_freshness: 0.15,
|
||||||
|
quota_remaining: 0.15,
|
||||||
|
latency: 0.10,
|
||||||
|
cost_lru: 0.10,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl PoolMemberScoreWeights {
|
||||||
|
pub fn normalized(self) -> Self {
|
||||||
|
let sanitized = Self {
|
||||||
|
manual_priority: finite_non_negative(self.manual_priority),
|
||||||
|
health: finite_non_negative(self.health),
|
||||||
|
probe_freshness: finite_non_negative(self.probe_freshness),
|
||||||
|
quota_remaining: finite_non_negative(self.quota_remaining),
|
||||||
|
latency: finite_non_negative(self.latency),
|
||||||
|
cost_lru: finite_non_negative(self.cost_lru),
|
||||||
|
};
|
||||||
|
let total = sanitized.manual_priority
|
||||||
|
+ sanitized.health
|
||||||
|
+ sanitized.probe_freshness
|
||||||
|
+ sanitized.quota_remaining
|
||||||
|
+ sanitized.latency
|
||||||
|
+ sanitized.cost_lru;
|
||||||
|
if total <= f64::EPSILON {
|
||||||
|
return sanitized;
|
||||||
|
}
|
||||||
|
Self {
|
||||||
|
manual_priority: sanitized.manual_priority / total,
|
||||||
|
health: sanitized.health / total,
|
||||||
|
probe_freshness: sanitized.probe_freshness / total,
|
||||||
|
quota_remaining: sanitized.quota_remaining / total,
|
||||||
|
latency: sanitized.latency / total,
|
||||||
|
cost_lru: sanitized.cost_lru / total,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn as_reason_json(self) -> Value {
|
||||||
|
json!({
|
||||||
|
"manual_priority": self.manual_priority,
|
||||||
|
"health": self.health,
|
||||||
|
"probe_freshness": self.probe_freshness,
|
||||||
|
"quota_remaining": self.quota_remaining,
|
||||||
|
"latency": self.latency,
|
||||||
|
"cost_lru": self.cost_lru
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Copy, PartialEq)]
|
||||||
|
pub struct PoolMemberScoreRules {
|
||||||
|
pub weights: PoolMemberScoreWeights,
|
||||||
|
pub probe_freshness_ttl_seconds: u64,
|
||||||
|
pub unschedulable_score_cap: f64,
|
||||||
|
pub probe_failure_penalty: f64,
|
||||||
|
pub request_failure_penalty: f64,
|
||||||
|
pub probe_failure_cooldown_threshold: u64,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for PoolMemberScoreRules {
|
||||||
|
fn default() -> Self {
|
||||||
|
Self {
|
||||||
|
weights: PoolMemberScoreWeights::default(),
|
||||||
|
probe_freshness_ttl_seconds: PROBE_FRESHNESS_TTL_SECONDS,
|
||||||
|
unschedulable_score_cap: UNSCHEDULABLE_SCORE_CAP,
|
||||||
|
probe_failure_penalty: PROBE_FAILURE_PENALTY,
|
||||||
|
request_failure_penalty: REQUEST_FAILURE_PENALTY,
|
||||||
|
probe_failure_cooldown_threshold: PROBE_FAILURE_COOLDOWN_THRESHOLD,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl PoolMemberScoreRules {
|
||||||
|
pub fn effective(self) -> Self {
|
||||||
|
let defaults = Self::default();
|
||||||
|
Self {
|
||||||
|
weights: self.weights.normalized(),
|
||||||
|
probe_freshness_ttl_seconds: if self.probe_freshness_ttl_seconds == 0 {
|
||||||
|
defaults.probe_freshness_ttl_seconds
|
||||||
|
} else {
|
||||||
|
self.probe_freshness_ttl_seconds
|
||||||
|
},
|
||||||
|
unschedulable_score_cap: if self.unschedulable_score_cap.is_finite() {
|
||||||
|
self.unschedulable_score_cap.clamp(0.0, 1.0)
|
||||||
|
} else {
|
||||||
|
defaults.unschedulable_score_cap
|
||||||
|
},
|
||||||
|
probe_failure_penalty: if self.probe_failure_penalty.is_finite() {
|
||||||
|
self.probe_failure_penalty.clamp(0.0, 1.0)
|
||||||
|
} else {
|
||||||
|
defaults.probe_failure_penalty
|
||||||
|
},
|
||||||
|
request_failure_penalty: if self.request_failure_penalty.is_finite() {
|
||||||
|
self.request_failure_penalty.clamp(0.0, 1.0)
|
||||||
|
} else {
|
||||||
|
defaults.request_failure_penalty
|
||||||
|
},
|
||||||
|
probe_failure_cooldown_threshold: self.probe_failure_cooldown_threshold,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct PoolMemberScoreInput {
|
||||||
|
pub identity: PoolMemberIdentity,
|
||||||
|
pub scope: PoolScoreScope,
|
||||||
|
pub internal_priority: i32,
|
||||||
|
pub is_active: bool,
|
||||||
|
pub health_score: Option<f64>,
|
||||||
|
pub quota_usage_ratio: Option<f64>,
|
||||||
|
pub quota_exhausted: bool,
|
||||||
|
pub account_blocked: bool,
|
||||||
|
pub oauth_invalid_reason: Option<String>,
|
||||||
|
pub circuit_open: bool,
|
||||||
|
pub success_count: u64,
|
||||||
|
pub error_count: u64,
|
||||||
|
pub total_response_time_ms: u64,
|
||||||
|
pub total_tokens: u64,
|
||||||
|
pub total_cost_usd: f64,
|
||||||
|
pub last_used_at: Option<u64>,
|
||||||
|
pub last_probe_success_at: Option<u64>,
|
||||||
|
pub probe_failure_count: u64,
|
||||||
|
pub probe_status: PoolMemberProbeStatus,
|
||||||
|
pub now_unix_secs: u64,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, PartialEq)]
|
||||||
|
pub struct PoolMemberScoreOutput {
|
||||||
|
pub score: f64,
|
||||||
|
pub hard_state: PoolMemberHardState,
|
||||||
|
pub score_reason: Value,
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn score_pool_member(input: &PoolMemberScoreInput) -> PoolMemberScoreOutput {
|
||||||
|
score_pool_member_with_rules(input, PoolMemberScoreRules::default())
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn score_pool_member_with_rules(
|
||||||
|
input: &PoolMemberScoreInput,
|
||||||
|
rules: PoolMemberScoreRules,
|
||||||
|
) -> PoolMemberScoreOutput {
|
||||||
|
let rules = rules.effective();
|
||||||
|
let weights = rules.weights;
|
||||||
|
let hard_state = derive_hard_state(input, &rules);
|
||||||
|
let manual_priority = manual_priority_score(input.internal_priority);
|
||||||
|
let health = input.health_score.unwrap_or(0.5).clamp(0.0, 1.0);
|
||||||
|
let probe_freshness = probe_freshness_score_with_ttl(
|
||||||
|
input.last_probe_success_at,
|
||||||
|
input.probe_status,
|
||||||
|
input.now_unix_secs,
|
||||||
|
rules.probe_freshness_ttl_seconds,
|
||||||
|
);
|
||||||
|
let quota_remaining = input
|
||||||
|
.quota_usage_ratio
|
||||||
|
.map(|ratio| 1.0 - ratio.clamp(0.0, 1.0))
|
||||||
|
.unwrap_or(0.5);
|
||||||
|
let latency = latency_score(input.success_count, input.total_response_time_ms);
|
||||||
|
let cost_lru = cost_lru_score(input.total_cost_usd, input.total_tokens, input.last_used_at);
|
||||||
|
|
||||||
|
let weighted_score = manual_priority * weights.manual_priority
|
||||||
|
+ health * weights.health
|
||||||
|
+ probe_freshness * weights.probe_freshness
|
||||||
|
+ quota_remaining * weights.quota_remaining
|
||||||
|
+ latency * weights.latency
|
||||||
|
+ cost_lru * weights.cost_lru;
|
||||||
|
let probe_failure_penalty =
|
||||||
|
(input.probe_failure_count.min(10) as f64 * rules.probe_failure_penalty).min(0.5);
|
||||||
|
let request_failure_penalty =
|
||||||
|
(input.error_count.min(20) as f64 * rules.request_failure_penalty).min(0.5);
|
||||||
|
let total_penalty = (probe_failure_penalty + request_failure_penalty).min(1.0);
|
||||||
|
let mut score = weighted_score - total_penalty;
|
||||||
|
if !hard_state.schedulable() {
|
||||||
|
score = score.min(rules.unschedulable_score_cap);
|
||||||
|
}
|
||||||
|
score = score.clamp(0.0, 1.0);
|
||||||
|
|
||||||
|
PoolMemberScoreOutput {
|
||||||
|
score,
|
||||||
|
hard_state,
|
||||||
|
score_reason: json!({
|
||||||
|
"weights": weights.as_reason_json(),
|
||||||
|
"factors": {
|
||||||
|
"manual_priority": manual_priority,
|
||||||
|
"health": health,
|
||||||
|
"probe_freshness": probe_freshness,
|
||||||
|
"quota_remaining": quota_remaining,
|
||||||
|
"latency": latency,
|
||||||
|
"cost_lru": cost_lru
|
||||||
|
},
|
||||||
|
"rules": {
|
||||||
|
"probe_freshness_ttl_seconds": rules.probe_freshness_ttl_seconds,
|
||||||
|
"unschedulable_score_cap": rules.unschedulable_score_cap,
|
||||||
|
"probe_failure_penalty": rules.probe_failure_penalty,
|
||||||
|
"request_failure_penalty": rules.request_failure_penalty,
|
||||||
|
"probe_failure_cooldown_threshold": rules.probe_failure_cooldown_threshold
|
||||||
|
},
|
||||||
|
"penalties": {
|
||||||
|
"probe_failure": probe_failure_penalty,
|
||||||
|
"request_failure": request_failure_penalty,
|
||||||
|
"total": total_penalty
|
||||||
|
},
|
||||||
|
"hard_state": hard_state.as_database(),
|
||||||
|
"score_version": POOL_SCORE_VERSION
|
||||||
|
}),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn finite_non_negative(value: f64) -> f64 {
|
||||||
|
if value.is_finite() {
|
||||||
|
value.max(0.0)
|
||||||
|
} else {
|
||||||
|
0.0
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn derive_hard_state(
|
||||||
|
input: &PoolMemberScoreInput,
|
||||||
|
rules: &PoolMemberScoreRules,
|
||||||
|
) -> PoolMemberHardState {
|
||||||
|
if !input.is_active {
|
||||||
|
return PoolMemberHardState::Inactive;
|
||||||
|
}
|
||||||
|
if let Some(reason) = input.oauth_invalid_reason.as_deref() {
|
||||||
|
let reason = reason.to_ascii_lowercase();
|
||||||
|
if reason.contains("ban") || reason.contains("blocked") || reason.contains("suspended") {
|
||||||
|
return PoolMemberHardState::Banned;
|
||||||
|
}
|
||||||
|
return PoolMemberHardState::AuthInvalid;
|
||||||
|
}
|
||||||
|
if input.account_blocked {
|
||||||
|
return PoolMemberHardState::Banned;
|
||||||
|
}
|
||||||
|
if input.quota_exhausted {
|
||||||
|
return PoolMemberHardState::QuotaExhausted;
|
||||||
|
}
|
||||||
|
if input.circuit_open {
|
||||||
|
return PoolMemberHardState::Cooldown;
|
||||||
|
}
|
||||||
|
if input.probe_status == PoolMemberProbeStatus::Failed
|
||||||
|
&& rules.probe_failure_cooldown_threshold > 0
|
||||||
|
&& input.probe_failure_count >= rules.probe_failure_cooldown_threshold
|
||||||
|
{
|
||||||
|
return PoolMemberHardState::Cooldown;
|
||||||
|
}
|
||||||
|
if input.health_score.is_some() || input.probe_status == PoolMemberProbeStatus::Ok {
|
||||||
|
PoolMemberHardState::Available
|
||||||
|
} else {
|
||||||
|
PoolMemberHardState::Unknown
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn manual_priority_score(internal_priority: i32) -> f64 {
|
||||||
|
(1.0 - (f64::from(internal_priority).clamp(0.0, 100.0) / 100.0)).clamp(0.0, 1.0)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn probe_freshness_score(
|
||||||
|
last_probe_success_at: Option<u64>,
|
||||||
|
probe_status: PoolMemberProbeStatus,
|
||||||
|
now_unix_secs: u64,
|
||||||
|
) -> f64 {
|
||||||
|
probe_freshness_score_with_ttl(
|
||||||
|
last_probe_success_at,
|
||||||
|
probe_status,
|
||||||
|
now_unix_secs,
|
||||||
|
PROBE_FRESHNESS_TTL_SECONDS,
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn probe_freshness_score_with_ttl(
|
||||||
|
last_probe_success_at: Option<u64>,
|
||||||
|
probe_status: PoolMemberProbeStatus,
|
||||||
|
now_unix_secs: u64,
|
||||||
|
ttl_seconds: u64,
|
||||||
|
) -> f64 {
|
||||||
|
if probe_status != PoolMemberProbeStatus::Ok {
|
||||||
|
return 0.0;
|
||||||
|
}
|
||||||
|
let Some(success_at) = last_probe_success_at else {
|
||||||
|
return 0.0;
|
||||||
|
};
|
||||||
|
let ttl_seconds = ttl_seconds.max(1);
|
||||||
|
let age = now_unix_secs.saturating_sub(success_at);
|
||||||
|
if age >= ttl_seconds {
|
||||||
|
0.0
|
||||||
|
} else {
|
||||||
|
1.0 - (age as f64 / ttl_seconds as f64)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn latency_score(success_count: u64, total_response_time_ms: u64) -> f64 {
|
||||||
|
if success_count == 0 || total_response_time_ms == 0 {
|
||||||
|
return 0.5;
|
||||||
|
}
|
||||||
|
let avg = total_response_time_ms as f64 / success_count as f64;
|
||||||
|
if avg <= 500.0 {
|
||||||
|
1.0
|
||||||
|
} else if avg >= 60_000.0 {
|
||||||
|
0.0
|
||||||
|
} else {
|
||||||
|
1.0 - ((avg - 500.0) / 59_500.0)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn cost_lru_score(total_cost_usd: f64, total_tokens: u64, last_used_at: Option<u64>) -> f64 {
|
||||||
|
let cost_penalty = if total_cost_usd.is_finite() {
|
||||||
|
(total_cost_usd.max(0.0) / 100.0).min(0.5)
|
||||||
|
} else {
|
||||||
|
0.0
|
||||||
|
};
|
||||||
|
let token_penalty = (total_tokens as f64 / 10_000_000.0).min(0.25);
|
||||||
|
let lru_bonus = if last_used_at.unwrap_or(0) == 0 {
|
||||||
|
0.25
|
||||||
|
} else {
|
||||||
|
0.0
|
||||||
|
};
|
||||||
|
(0.75 - cost_penalty - token_penalty + lru_bonus).clamp(0.0, 1.0)
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
use aether_data_contracts::repository::pool_scores::{
|
||||||
|
POOL_KIND_PROVIDER_KEY_POOL, POOL_MEMBER_KIND_PROVIDER_API_KEY, POOL_SCORE_SCOPE_KIND_MODEL,
|
||||||
|
};
|
||||||
|
|
||||||
|
fn input() -> PoolMemberScoreInput {
|
||||||
|
PoolMemberScoreInput {
|
||||||
|
identity: PoolMemberIdentity {
|
||||||
|
pool_kind: POOL_KIND_PROVIDER_KEY_POOL.to_string(),
|
||||||
|
pool_id: "provider-1".to_string(),
|
||||||
|
member_kind: POOL_MEMBER_KIND_PROVIDER_API_KEY.to_string(),
|
||||||
|
member_id: "key-1".to_string(),
|
||||||
|
},
|
||||||
|
scope: PoolScoreScope {
|
||||||
|
capability: "openai:responses".to_string(),
|
||||||
|
scope_kind: POOL_SCORE_SCOPE_KIND_MODEL.to_string(),
|
||||||
|
scope_id: Some("model-1".to_string()),
|
||||||
|
},
|
||||||
|
internal_priority: 10,
|
||||||
|
is_active: true,
|
||||||
|
health_score: Some(1.0),
|
||||||
|
quota_usage_ratio: Some(0.1),
|
||||||
|
quota_exhausted: false,
|
||||||
|
account_blocked: false,
|
||||||
|
oauth_invalid_reason: None,
|
||||||
|
circuit_open: false,
|
||||||
|
success_count: 10,
|
||||||
|
error_count: 0,
|
||||||
|
total_response_time_ms: 2_000,
|
||||||
|
total_tokens: 10,
|
||||||
|
total_cost_usd: 0.01,
|
||||||
|
last_used_at: None,
|
||||||
|
last_probe_success_at: Some(1_000),
|
||||||
|
probe_failure_count: 0,
|
||||||
|
probe_status: PoolMemberProbeStatus::Ok,
|
||||||
|
now_unix_secs: 1_000,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn hard_state_caps_unavailable_member_score() {
|
||||||
|
let mut input = input();
|
||||||
|
input.oauth_invalid_reason = Some("token invalid".to_string());
|
||||||
|
|
||||||
|
let output = score_pool_member(&input);
|
||||||
|
|
||||||
|
assert_eq!(output.hard_state, PoolMemberHardState::AuthInvalid);
|
||||||
|
assert!(output.score <= 0.05);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn probe_freshness_has_ttl() {
|
||||||
|
assert_eq!(
|
||||||
|
probe_freshness_score(Some(1_000), PoolMemberProbeStatus::Ok, 1_000),
|
||||||
|
1.0
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
probe_freshness_score(
|
||||||
|
Some(1_000),
|
||||||
|
PoolMemberProbeStatus::Ok,
|
||||||
|
1_000 + PROBE_FRESHNESS_TTL_SECONDS
|
||||||
|
),
|
||||||
|
0.0
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn custom_rules_change_weights_and_probe_ttl() {
|
||||||
|
let mut input = input();
|
||||||
|
input.last_probe_success_at = Some(1_000);
|
||||||
|
input.now_unix_secs = 1_900;
|
||||||
|
let rules = PoolMemberScoreRules {
|
||||||
|
weights: PoolMemberScoreWeights {
|
||||||
|
manual_priority: 0.0,
|
||||||
|
health: 0.0,
|
||||||
|
probe_freshness: 1.0,
|
||||||
|
quota_remaining: 0.0,
|
||||||
|
latency: 0.0,
|
||||||
|
cost_lru: 0.0,
|
||||||
|
},
|
||||||
|
probe_freshness_ttl_seconds: 1_000,
|
||||||
|
unschedulable_score_cap: 0.05,
|
||||||
|
probe_failure_penalty: 0.0,
|
||||||
|
request_failure_penalty: 0.0,
|
||||||
|
probe_failure_cooldown_threshold: PROBE_FAILURE_COOLDOWN_THRESHOLD,
|
||||||
|
};
|
||||||
|
|
||||||
|
let output = score_pool_member_with_rules(&input, rules);
|
||||||
|
|
||||||
|
assert!((output.score - 0.1).abs() < 0.000_001);
|
||||||
|
assert_eq!(
|
||||||
|
output.score_reason["rules"]["probe_freshness_ttl_seconds"],
|
||||||
|
1_000
|
||||||
|
);
|
||||||
|
assert_eq!(output.score_reason["weights"]["probe_freshness"], 1.0);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn custom_rules_normalize_weights() {
|
||||||
|
let rules = PoolMemberScoreRules {
|
||||||
|
weights: PoolMemberScoreWeights {
|
||||||
|
manual_priority: 2.0,
|
||||||
|
health: 2.0,
|
||||||
|
probe_freshness: 0.0,
|
||||||
|
quota_remaining: 0.0,
|
||||||
|
latency: -1.0,
|
||||||
|
cost_lru: f64::NAN,
|
||||||
|
},
|
||||||
|
probe_freshness_ttl_seconds: 0,
|
||||||
|
unschedulable_score_cap: f64::INFINITY,
|
||||||
|
probe_failure_penalty: f64::NAN,
|
||||||
|
request_failure_penalty: f64::INFINITY,
|
||||||
|
probe_failure_cooldown_threshold: 2,
|
||||||
|
}
|
||||||
|
.effective();
|
||||||
|
|
||||||
|
assert_eq!(rules.weights.manual_priority, 0.5);
|
||||||
|
assert_eq!(rules.weights.health, 0.5);
|
||||||
|
assert_eq!(
|
||||||
|
rules.probe_freshness_ttl_seconds,
|
||||||
|
PROBE_FRESHNESS_TTL_SECONDS
|
||||||
|
);
|
||||||
|
assert_eq!(rules.unschedulable_score_cap, UNSCHEDULABLE_SCORE_CAP);
|
||||||
|
assert_eq!(rules.probe_failure_penalty, PROBE_FAILURE_PENALTY);
|
||||||
|
assert_eq!(rules.request_failure_penalty, REQUEST_FAILURE_PENALTY);
|
||||||
|
assert_eq!(rules.probe_failure_cooldown_threshold, 2);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn custom_rules_preserve_zero_weight_total() {
|
||||||
|
let rules = PoolMemberScoreRules {
|
||||||
|
weights: PoolMemberScoreWeights {
|
||||||
|
manual_priority: 0.0,
|
||||||
|
health: 0.0,
|
||||||
|
probe_freshness: 0.0,
|
||||||
|
quota_remaining: 0.0,
|
||||||
|
latency: 0.0,
|
||||||
|
cost_lru: 0.0,
|
||||||
|
},
|
||||||
|
..PoolMemberScoreRules::default()
|
||||||
|
}
|
||||||
|
.effective();
|
||||||
|
|
||||||
|
assert_eq!(
|
||||||
|
rules.weights,
|
||||||
|
PoolMemberScoreWeights {
|
||||||
|
manual_priority: 0.0,
|
||||||
|
health: 0.0,
|
||||||
|
probe_freshness: 0.0,
|
||||||
|
quota_remaining: 0.0,
|
||||||
|
latency: 0.0,
|
||||||
|
cost_lru: 0.0,
|
||||||
|
}
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn repeated_probe_failures_move_member_to_cooldown() {
|
||||||
|
let mut input = input();
|
||||||
|
input.probe_status = PoolMemberProbeStatus::Failed;
|
||||||
|
input.probe_failure_count = 3;
|
||||||
|
|
||||||
|
let output = score_pool_member(&input);
|
||||||
|
|
||||||
|
assert_eq!(output.hard_state, PoolMemberHardState::Cooldown);
|
||||||
|
assert!(output.score <= UNSCHEDULABLE_SCORE_CAP);
|
||||||
|
assert_eq!(
|
||||||
|
output.score_reason["rules"]["probe_failure_cooldown_threshold"],
|
||||||
|
PROBE_FAILURE_COOLDOWN_THRESHOLD
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn probe_and_request_failures_penalize_schedulable_score() {
|
||||||
|
let mut input = input();
|
||||||
|
input.last_probe_success_at = None;
|
||||||
|
input.probe_status = PoolMemberProbeStatus::Never;
|
||||||
|
input.probe_failure_count = 1;
|
||||||
|
input.error_count = 2;
|
||||||
|
let rules = PoolMemberScoreRules {
|
||||||
|
probe_failure_penalty: 0.1,
|
||||||
|
request_failure_penalty: 0.01,
|
||||||
|
probe_failure_cooldown_threshold: 3,
|
||||||
|
..PoolMemberScoreRules::default()
|
||||||
|
};
|
||||||
|
|
||||||
|
let output = score_pool_member_with_rules(&input, rules);
|
||||||
|
|
||||||
|
assert_eq!(output.hard_state, PoolMemberHardState::Available);
|
||||||
|
assert_eq!(output.score_reason["penalties"]["probe_failure"], 0.1);
|
||||||
|
assert_eq!(output.score_reason["penalties"]["request_failure"], 0.02);
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -3,6 +3,6 @@ mod types;
|
|||||||
pub use types::{
|
pub use types::{
|
||||||
MinimalCandidateSelectionReadRepository, MinimalCandidateSelectionRepository,
|
MinimalCandidateSelectionReadRepository, MinimalCandidateSelectionRepository,
|
||||||
StoredMinimalCandidateSelectionRow, StoredPoolKeyCandidateOrder,
|
StoredMinimalCandidateSelectionRow, StoredPoolKeyCandidateOrder,
|
||||||
StoredPoolKeyCandidateRowsQuery, StoredProviderModelMapping,
|
StoredPoolKeyCandidateRowsByKeyIdsQuery, StoredPoolKeyCandidateRowsQuery,
|
||||||
StoredRequestedModelCandidateRowsQuery,
|
StoredProviderModelMapping, StoredRequestedModelCandidateRowsQuery,
|
||||||
};
|
};
|
||||||
|
|||||||
@@ -67,6 +67,16 @@ pub struct StoredPoolKeyCandidateRowsQuery {
|
|||||||
pub limit: u32,
|
pub limit: u32,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||||
|
pub struct StoredPoolKeyCandidateRowsByKeyIdsQuery {
|
||||||
|
pub api_format: String,
|
||||||
|
pub provider_id: String,
|
||||||
|
pub endpoint_id: String,
|
||||||
|
pub model_id: String,
|
||||||
|
pub selected_provider_model_name: String,
|
||||||
|
pub key_ids: Vec<String>,
|
||||||
|
}
|
||||||
|
|
||||||
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||||
pub struct StoredRequestedModelCandidateRowsQuery {
|
pub struct StoredRequestedModelCandidateRowsQuery {
|
||||||
pub api_format: String,
|
pub api_format: String,
|
||||||
@@ -124,6 +134,11 @@ pub trait MinimalCandidateSelectionReadRepository: Send + Sync {
|
|||||||
&self,
|
&self,
|
||||||
query: &StoredPoolKeyCandidateRowsQuery,
|
query: &StoredPoolKeyCandidateRowsQuery,
|
||||||
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, crate::DataLayerError>;
|
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, crate::DataLayerError>;
|
||||||
|
|
||||||
|
async fn list_pool_key_rows_for_group_key_ids(
|
||||||
|
&self,
|
||||||
|
query: &StoredPoolKeyCandidateRowsByKeyIdsQuery,
|
||||||
|
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, crate::DataLayerError>;
|
||||||
}
|
}
|
||||||
|
|
||||||
pub trait MinimalCandidateSelectionRepository:
|
pub trait MinimalCandidateSelectionRepository:
|
||||||
|
|||||||
@@ -3,6 +3,7 @@ pub mod billing;
|
|||||||
pub mod candidate_selection;
|
pub mod candidate_selection;
|
||||||
pub mod candidates;
|
pub mod candidates;
|
||||||
pub mod global_models;
|
pub mod global_models;
|
||||||
|
pub mod pool_scores;
|
||||||
pub mod provider_catalog;
|
pub mod provider_catalog;
|
||||||
pub mod quota;
|
pub mod quota;
|
||||||
pub mod settlement;
|
pub mod settlement;
|
||||||
|
|||||||
@@ -0,0 +1,11 @@
|
|||||||
|
mod types;
|
||||||
|
|
||||||
|
pub use types::{
|
||||||
|
GetPoolMemberScoresByIdsQuery, ListPoolMemberProbeCandidatesQuery, ListPoolMemberScoresQuery,
|
||||||
|
ListRankedPoolMembersQuery, PoolMemberHardState, PoolMemberIdentity, PoolMemberProbeAttempt,
|
||||||
|
PoolMemberProbeResult, PoolMemberProbeStatus, PoolMemberScheduleFeedback,
|
||||||
|
PoolMemberScoreRepository, PoolMemberScoreWriteRepository, PoolScoreReadRepository,
|
||||||
|
PoolScoreScope, StoredPoolMemberScore, UpsertPoolMemberScore, POOL_KIND_PROVIDER_KEY_POOL,
|
||||||
|
POOL_MEMBER_KIND_PROVIDER_API_KEY, POOL_SCORE_CAPABILITY_API_FORMAT,
|
||||||
|
POOL_SCORE_SCOPE_KIND_MODEL,
|
||||||
|
};
|
||||||
359
crates/aether-data-contracts/src/repository/pool_scores/types.rs
Normal file
359
crates/aether-data-contracts/src/repository/pool_scores/types.rs
Normal file
@@ -0,0 +1,359 @@
|
|||||||
|
use async_trait::async_trait;
|
||||||
|
|
||||||
|
pub const POOL_KIND_PROVIDER_KEY_POOL: &str = "provider_key_pool";
|
||||||
|
pub const POOL_MEMBER_KIND_PROVIDER_API_KEY: &str = "provider_api_key";
|
||||||
|
pub const POOL_SCORE_CAPABILITY_API_FORMAT: &str = "api_format";
|
||||||
|
pub const POOL_SCORE_SCOPE_KIND_MODEL: &str = "model";
|
||||||
|
|
||||||
|
#[derive(
|
||||||
|
Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, serde::Serialize, serde::Deserialize,
|
||||||
|
)]
|
||||||
|
#[serde(rename_all = "snake_case")]
|
||||||
|
pub enum PoolMemberHardState {
|
||||||
|
Available,
|
||||||
|
Unknown,
|
||||||
|
Cooldown,
|
||||||
|
QuotaExhausted,
|
||||||
|
AuthInvalid,
|
||||||
|
Banned,
|
||||||
|
Inactive,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl PoolMemberHardState {
|
||||||
|
pub fn as_database(self) -> &'static str {
|
||||||
|
match self {
|
||||||
|
Self::Available => "available",
|
||||||
|
Self::Unknown => "unknown",
|
||||||
|
Self::Cooldown => "cooldown",
|
||||||
|
Self::QuotaExhausted => "quota_exhausted",
|
||||||
|
Self::AuthInvalid => "auth_invalid",
|
||||||
|
Self::Banned => "banned",
|
||||||
|
Self::Inactive => "inactive",
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn from_database(value: &str) -> Result<Self, crate::DataLayerError> {
|
||||||
|
match value.trim().to_ascii_lowercase().as_str() {
|
||||||
|
"available" => Ok(Self::Available),
|
||||||
|
"unknown" => Ok(Self::Unknown),
|
||||||
|
"cooldown" => Ok(Self::Cooldown),
|
||||||
|
"quota_exhausted" => Ok(Self::QuotaExhausted),
|
||||||
|
"auth_invalid" => Ok(Self::AuthInvalid),
|
||||||
|
"banned" => Ok(Self::Banned),
|
||||||
|
"inactive" => Ok(Self::Inactive),
|
||||||
|
other => Err(crate::DataLayerError::UnexpectedValue(format!(
|
||||||
|
"unknown pool member hard_state: {other}"
|
||||||
|
))),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn schedulable(self) -> bool {
|
||||||
|
matches!(self, Self::Available | Self::Unknown)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(
|
||||||
|
Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, serde::Serialize, serde::Deserialize,
|
||||||
|
)]
|
||||||
|
#[serde(rename_all = "snake_case")]
|
||||||
|
pub enum PoolMemberProbeStatus {
|
||||||
|
Never,
|
||||||
|
Ok,
|
||||||
|
Failed,
|
||||||
|
Stale,
|
||||||
|
InProgress,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl PoolMemberProbeStatus {
|
||||||
|
pub fn as_database(self) -> &'static str {
|
||||||
|
match self {
|
||||||
|
Self::Never => "never",
|
||||||
|
Self::Ok => "ok",
|
||||||
|
Self::Failed => "failed",
|
||||||
|
Self::Stale => "stale",
|
||||||
|
Self::InProgress => "in_progress",
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn from_database(value: &str) -> Result<Self, crate::DataLayerError> {
|
||||||
|
match value.trim().to_ascii_lowercase().as_str() {
|
||||||
|
"never" => Ok(Self::Never),
|
||||||
|
"ok" => Ok(Self::Ok),
|
||||||
|
"failed" => Ok(Self::Failed),
|
||||||
|
"stale" => Ok(Self::Stale),
|
||||||
|
"in_progress" => Ok(Self::InProgress),
|
||||||
|
other => Err(crate::DataLayerError::UnexpectedValue(format!(
|
||||||
|
"unknown pool member probe_status: {other}"
|
||||||
|
))),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||||
|
pub struct PoolScoreScope {
|
||||||
|
pub capability: String,
|
||||||
|
pub scope_kind: String,
|
||||||
|
pub scope_id: Option<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||||
|
pub struct PoolMemberIdentity {
|
||||||
|
pub pool_kind: String,
|
||||||
|
pub pool_id: String,
|
||||||
|
pub member_kind: String,
|
||||||
|
pub member_id: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl PoolMemberIdentity {
|
||||||
|
pub fn provider_api_key(provider_id: impl Into<String>, key_id: impl Into<String>) -> Self {
|
||||||
|
Self {
|
||||||
|
pool_kind: POOL_KIND_PROVIDER_KEY_POOL.to_string(),
|
||||||
|
pool_id: provider_id.into(),
|
||||||
|
member_kind: POOL_MEMBER_KIND_PROVIDER_API_KEY.to_string(),
|
||||||
|
member_id: key_id.into(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||||
|
pub struct StoredPoolMemberScore {
|
||||||
|
pub id: String,
|
||||||
|
pub pool_kind: String,
|
||||||
|
pub pool_id: String,
|
||||||
|
pub member_kind: String,
|
||||||
|
pub member_id: String,
|
||||||
|
pub capability: String,
|
||||||
|
pub scope_kind: String,
|
||||||
|
pub scope_id: Option<String>,
|
||||||
|
pub score: f64,
|
||||||
|
pub hard_state: PoolMemberHardState,
|
||||||
|
pub score_version: u64,
|
||||||
|
pub score_reason: serde_json::Value,
|
||||||
|
pub last_ranked_at: Option<u64>,
|
||||||
|
pub last_scheduled_at: Option<u64>,
|
||||||
|
pub last_success_at: Option<u64>,
|
||||||
|
pub last_failure_at: Option<u64>,
|
||||||
|
pub failure_count: u64,
|
||||||
|
pub last_probe_attempt_at: Option<u64>,
|
||||||
|
pub last_probe_success_at: Option<u64>,
|
||||||
|
pub last_probe_failure_at: Option<u64>,
|
||||||
|
pub probe_failure_count: u64,
|
||||||
|
pub probe_status: PoolMemberProbeStatus,
|
||||||
|
pub updated_at: u64,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||||
|
pub struct UpsertPoolMemberScore {
|
||||||
|
pub id: String,
|
||||||
|
pub identity: PoolMemberIdentity,
|
||||||
|
pub scope: PoolScoreScope,
|
||||||
|
pub score: f64,
|
||||||
|
pub hard_state: PoolMemberHardState,
|
||||||
|
pub score_version: u64,
|
||||||
|
pub score_reason: serde_json::Value,
|
||||||
|
pub last_ranked_at: Option<u64>,
|
||||||
|
pub last_scheduled_at: Option<u64>,
|
||||||
|
pub last_success_at: Option<u64>,
|
||||||
|
pub last_failure_at: Option<u64>,
|
||||||
|
pub failure_count: u64,
|
||||||
|
pub last_probe_attempt_at: Option<u64>,
|
||||||
|
pub last_probe_success_at: Option<u64>,
|
||||||
|
pub last_probe_failure_at: Option<u64>,
|
||||||
|
pub probe_failure_count: u64,
|
||||||
|
pub probe_status: PoolMemberProbeStatus,
|
||||||
|
pub updated_at: u64,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl UpsertPoolMemberScore {
|
||||||
|
pub fn validate(&self) -> Result<(), crate::DataLayerError> {
|
||||||
|
validate_non_empty(&self.id, "pool_member_scores.id")?;
|
||||||
|
validate_non_empty(&self.identity.pool_kind, "pool_member_scores.pool_kind")?;
|
||||||
|
validate_non_empty(&self.identity.pool_id, "pool_member_scores.pool_id")?;
|
||||||
|
validate_non_empty(&self.identity.member_kind, "pool_member_scores.member_kind")?;
|
||||||
|
validate_non_empty(&self.identity.member_id, "pool_member_scores.member_id")?;
|
||||||
|
validate_non_empty(&self.scope.capability, "pool_member_scores.capability")?;
|
||||||
|
validate_non_empty(&self.scope.scope_kind, "pool_member_scores.scope_kind")?;
|
||||||
|
if !self.score.is_finite() {
|
||||||
|
return Err(crate::DataLayerError::InvalidInput(
|
||||||
|
"pool_member_scores.score must be finite".to_string(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn into_stored(self) -> StoredPoolMemberScore {
|
||||||
|
StoredPoolMemberScore {
|
||||||
|
id: self.id,
|
||||||
|
pool_kind: self.identity.pool_kind,
|
||||||
|
pool_id: self.identity.pool_id,
|
||||||
|
member_kind: self.identity.member_kind,
|
||||||
|
member_id: self.identity.member_id,
|
||||||
|
capability: self.scope.capability,
|
||||||
|
scope_kind: self.scope.scope_kind,
|
||||||
|
scope_id: self.scope.scope_id,
|
||||||
|
score: self.score,
|
||||||
|
hard_state: self.hard_state,
|
||||||
|
score_version: self.score_version,
|
||||||
|
score_reason: self.score_reason,
|
||||||
|
last_ranked_at: self.last_ranked_at,
|
||||||
|
last_scheduled_at: self.last_scheduled_at,
|
||||||
|
last_success_at: self.last_success_at,
|
||||||
|
last_failure_at: self.last_failure_at,
|
||||||
|
failure_count: self.failure_count,
|
||||||
|
last_probe_attempt_at: self.last_probe_attempt_at,
|
||||||
|
last_probe_success_at: self.last_probe_success_at,
|
||||||
|
last_probe_failure_at: self.last_probe_failure_at,
|
||||||
|
probe_failure_count: self.probe_failure_count,
|
||||||
|
probe_status: self.probe_status,
|
||||||
|
updated_at: self.updated_at,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||||
|
pub struct ListRankedPoolMembersQuery {
|
||||||
|
pub pool_kind: String,
|
||||||
|
pub pool_id: String,
|
||||||
|
pub capability: String,
|
||||||
|
pub scope_kind: String,
|
||||||
|
pub scope_id: Option<String>,
|
||||||
|
pub hard_states: Vec<PoolMemberHardState>,
|
||||||
|
pub probe_statuses: Option<Vec<PoolMemberProbeStatus>>,
|
||||||
|
pub offset: usize,
|
||||||
|
pub limit: usize,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||||
|
pub struct ListPoolMemberScoresQuery {
|
||||||
|
pub pool_kind: String,
|
||||||
|
pub pool_id: String,
|
||||||
|
pub capability: Option<String>,
|
||||||
|
pub scope_kind: Option<String>,
|
||||||
|
pub scope_id: Option<String>,
|
||||||
|
pub hard_states: Vec<PoolMemberHardState>,
|
||||||
|
pub probe_statuses: Option<Vec<PoolMemberProbeStatus>>,
|
||||||
|
pub offset: usize,
|
||||||
|
pub limit: usize,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||||
|
pub struct ListPoolMemberProbeCandidatesQuery {
|
||||||
|
pub pool_kind: String,
|
||||||
|
pub pool_id: String,
|
||||||
|
pub capability: Option<String>,
|
||||||
|
pub stale_before_unix_secs: u64,
|
||||||
|
pub limit: usize,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||||
|
pub struct GetPoolMemberScoresByIdsQuery {
|
||||||
|
pub ids: Vec<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||||
|
pub struct PoolMemberProbeResult {
|
||||||
|
pub identity: PoolMemberIdentity,
|
||||||
|
pub scope: Option<PoolScoreScope>,
|
||||||
|
pub attempted_at: u64,
|
||||||
|
pub succeeded: bool,
|
||||||
|
pub hard_state: Option<PoolMemberHardState>,
|
||||||
|
pub probe_status: PoolMemberProbeStatus,
|
||||||
|
pub score_reason_patch: Option<serde_json::Value>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||||
|
pub struct PoolMemberProbeAttempt {
|
||||||
|
pub identity: PoolMemberIdentity,
|
||||||
|
pub scope: Option<PoolScoreScope>,
|
||||||
|
pub attempted_at: u64,
|
||||||
|
pub score_reason_patch: Option<serde_json::Value>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||||
|
pub struct PoolMemberScheduleFeedback {
|
||||||
|
pub identity: PoolMemberIdentity,
|
||||||
|
pub scope: Option<PoolScoreScope>,
|
||||||
|
pub scheduled_at: u64,
|
||||||
|
pub succeeded: Option<bool>,
|
||||||
|
pub hard_state: Option<PoolMemberHardState>,
|
||||||
|
pub score_delta: Option<i32>,
|
||||||
|
pub score_reason_patch: Option<serde_json::Value>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
pub trait PoolScoreReadRepository: Send + Sync {
|
||||||
|
async fn list_ranked_pool_members(
|
||||||
|
&self,
|
||||||
|
query: &ListRankedPoolMembersQuery,
|
||||||
|
) -> Result<Vec<StoredPoolMemberScore>, crate::DataLayerError>;
|
||||||
|
|
||||||
|
async fn list_pool_member_scores(
|
||||||
|
&self,
|
||||||
|
query: &ListPoolMemberScoresQuery,
|
||||||
|
) -> Result<Vec<StoredPoolMemberScore>, crate::DataLayerError>;
|
||||||
|
|
||||||
|
async fn list_pool_member_probe_candidates(
|
||||||
|
&self,
|
||||||
|
query: &ListPoolMemberProbeCandidatesQuery,
|
||||||
|
) -> Result<Vec<StoredPoolMemberScore>, crate::DataLayerError>;
|
||||||
|
|
||||||
|
async fn get_pool_member_scores_by_ids(
|
||||||
|
&self,
|
||||||
|
query: &GetPoolMemberScoresByIdsQuery,
|
||||||
|
) -> Result<Vec<StoredPoolMemberScore>, crate::DataLayerError>;
|
||||||
|
}
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
pub trait PoolMemberScoreWriteRepository: Send + Sync {
|
||||||
|
async fn upsert_pool_member_score(
|
||||||
|
&self,
|
||||||
|
score: UpsertPoolMemberScore,
|
||||||
|
) -> Result<StoredPoolMemberScore, crate::DataLayerError>;
|
||||||
|
|
||||||
|
async fn mark_pool_member_probe_in_progress(
|
||||||
|
&self,
|
||||||
|
attempt: PoolMemberProbeAttempt,
|
||||||
|
) -> Result<usize, crate::DataLayerError>;
|
||||||
|
|
||||||
|
async fn record_pool_member_probe_result(
|
||||||
|
&self,
|
||||||
|
result: PoolMemberProbeResult,
|
||||||
|
) -> Result<usize, crate::DataLayerError>;
|
||||||
|
|
||||||
|
async fn record_pool_member_schedule_feedback(
|
||||||
|
&self,
|
||||||
|
feedback: PoolMemberScheduleFeedback,
|
||||||
|
) -> Result<usize, crate::DataLayerError>;
|
||||||
|
|
||||||
|
async fn mark_pool_member_hard_state(
|
||||||
|
&self,
|
||||||
|
identity: &PoolMemberIdentity,
|
||||||
|
scope: Option<&PoolScoreScope>,
|
||||||
|
hard_state: PoolMemberHardState,
|
||||||
|
updated_at: u64,
|
||||||
|
) -> Result<usize, crate::DataLayerError>;
|
||||||
|
|
||||||
|
async fn delete_pool_member_scores_for_member(
|
||||||
|
&self,
|
||||||
|
identity: &PoolMemberIdentity,
|
||||||
|
) -> Result<usize, crate::DataLayerError>;
|
||||||
|
}
|
||||||
|
|
||||||
|
pub trait PoolMemberScoreRepository:
|
||||||
|
PoolScoreReadRepository + PoolMemberScoreWriteRepository + Send + Sync
|
||||||
|
{
|
||||||
|
}
|
||||||
|
|
||||||
|
impl<T> PoolMemberScoreRepository for T where
|
||||||
|
T: PoolScoreReadRepository + PoolMemberScoreWriteRepository + Send + Sync
|
||||||
|
{
|
||||||
|
}
|
||||||
|
|
||||||
|
fn validate_non_empty(value: &str, field: &str) -> Result<(), crate::DataLayerError> {
|
||||||
|
if value.trim().is_empty() {
|
||||||
|
return Err(crate::DataLayerError::InvalidInput(format!(
|
||||||
|
"{field} is empty"
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
@@ -294,6 +294,36 @@ CREATE TABLE IF NOT EXISTS provider_api_keys (
|
|||||||
KEY provider_api_keys_provider_id_idx (provider_id)
|
KEY provider_api_keys_provider_id_idx (provider_id)
|
||||||
);
|
);
|
||||||
|
|
||||||
|
CREATE TABLE IF NOT EXISTS pool_member_scores (
|
||||||
|
id VARCHAR(192) PRIMARY KEY,
|
||||||
|
pool_kind VARCHAR(64) NOT NULL,
|
||||||
|
pool_id VARCHAR(64) NOT NULL,
|
||||||
|
member_kind VARCHAR(64) NOT NULL,
|
||||||
|
member_id VARCHAR(64) NOT NULL,
|
||||||
|
capability VARCHAR(64) NOT NULL,
|
||||||
|
scope_kind VARCHAR(64) NOT NULL,
|
||||||
|
scope_id VARCHAR(128),
|
||||||
|
score DOUBLE NOT NULL DEFAULT 0,
|
||||||
|
hard_state VARCHAR(64) NOT NULL DEFAULT 'unknown',
|
||||||
|
score_version BIGINT NOT NULL DEFAULT 1,
|
||||||
|
score_reason TEXT NOT NULL,
|
||||||
|
last_ranked_at BIGINT,
|
||||||
|
last_scheduled_at BIGINT,
|
||||||
|
last_success_at BIGINT,
|
||||||
|
last_failure_at BIGINT,
|
||||||
|
failure_count BIGINT NOT NULL DEFAULT 0,
|
||||||
|
last_probe_attempt_at BIGINT,
|
||||||
|
last_probe_success_at BIGINT,
|
||||||
|
last_probe_failure_at BIGINT,
|
||||||
|
probe_failure_count BIGINT NOT NULL DEFAULT 0,
|
||||||
|
probe_status VARCHAR(64) NOT NULL DEFAULT 'never',
|
||||||
|
updated_at BIGINT NOT NULL,
|
||||||
|
KEY pool_member_scores_rank_idx (pool_kind, pool_id, capability, scope_kind, scope_id, hard_state, score DESC),
|
||||||
|
KEY pool_member_scores_member_idx (pool_kind, pool_id, member_kind, member_id),
|
||||||
|
KEY pool_member_scores_probe_idx (pool_kind, pool_id, probe_status, last_probe_success_at),
|
||||||
|
KEY pool_member_scores_updated_at_idx (updated_at)
|
||||||
|
);
|
||||||
|
|
||||||
CREATE TABLE IF NOT EXISTS gemini_file_mappings (
|
CREATE TABLE IF NOT EXISTS gemini_file_mappings (
|
||||||
id VARCHAR(64) PRIMARY KEY,
|
id VARCHAR(64) PRIMARY KEY,
|
||||||
file_name VARCHAR(512) NOT NULL,
|
file_name VARCHAR(512) NOT NULL,
|
||||||
|
|||||||
@@ -0,0 +1,29 @@
|
|||||||
|
CREATE TABLE IF NOT EXISTS pool_member_scores (
|
||||||
|
id VARCHAR(192) PRIMARY KEY,
|
||||||
|
pool_kind VARCHAR(64) NOT NULL,
|
||||||
|
pool_id VARCHAR(64) NOT NULL,
|
||||||
|
member_kind VARCHAR(64) NOT NULL,
|
||||||
|
member_id VARCHAR(64) NOT NULL,
|
||||||
|
capability VARCHAR(64) NOT NULL,
|
||||||
|
scope_kind VARCHAR(64) NOT NULL,
|
||||||
|
scope_id VARCHAR(128),
|
||||||
|
score DOUBLE NOT NULL DEFAULT 0,
|
||||||
|
hard_state VARCHAR(64) NOT NULL DEFAULT 'unknown',
|
||||||
|
score_version BIGINT NOT NULL DEFAULT 1,
|
||||||
|
score_reason TEXT NOT NULL,
|
||||||
|
last_ranked_at BIGINT,
|
||||||
|
last_scheduled_at BIGINT,
|
||||||
|
last_success_at BIGINT,
|
||||||
|
last_failure_at BIGINT,
|
||||||
|
failure_count BIGINT NOT NULL DEFAULT 0,
|
||||||
|
last_probe_attempt_at BIGINT,
|
||||||
|
last_probe_success_at BIGINT,
|
||||||
|
last_probe_failure_at BIGINT,
|
||||||
|
probe_failure_count BIGINT NOT NULL DEFAULT 0,
|
||||||
|
probe_status VARCHAR(64) NOT NULL DEFAULT 'never',
|
||||||
|
updated_at BIGINT NOT NULL,
|
||||||
|
KEY pool_member_scores_rank_idx (pool_kind, pool_id, capability, scope_kind, scope_id, hard_state, score DESC),
|
||||||
|
KEY pool_member_scores_member_idx (pool_kind, pool_id, member_kind, member_id),
|
||||||
|
KEY pool_member_scores_probe_idx (pool_kind, pool_id, probe_status, last_probe_success_at),
|
||||||
|
KEY pool_member_scores_updated_at_idx (updated_at)
|
||||||
|
);
|
||||||
@@ -534,6 +534,38 @@ CREATE TABLE IF NOT EXISTS public.provider_api_keys (
|
|||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
--
|
||||||
|
-- Name: pool_member_scores; Type: TABLE; Schema: public; Owner: -
|
||||||
|
--
|
||||||
|
|
||||||
|
CREATE TABLE IF NOT EXISTS public.pool_member_scores (
|
||||||
|
id character varying(192) NOT NULL,
|
||||||
|
pool_kind character varying(64) NOT NULL,
|
||||||
|
pool_id character varying(64) NOT NULL,
|
||||||
|
member_kind character varying(64) NOT NULL,
|
||||||
|
member_id character varying(64) NOT NULL,
|
||||||
|
capability character varying(64) NOT NULL,
|
||||||
|
scope_kind character varying(64) NOT NULL,
|
||||||
|
scope_id character varying(128),
|
||||||
|
score double precision DEFAULT 0 NOT NULL,
|
||||||
|
hard_state character varying(64) DEFAULT 'unknown'::character varying NOT NULL,
|
||||||
|
score_version bigint DEFAULT 1 NOT NULL,
|
||||||
|
score_reason jsonb NOT NULL,
|
||||||
|
last_ranked_at bigint,
|
||||||
|
last_scheduled_at bigint,
|
||||||
|
last_success_at bigint,
|
||||||
|
last_failure_at bigint,
|
||||||
|
failure_count bigint DEFAULT 0 NOT NULL,
|
||||||
|
last_probe_attempt_at bigint,
|
||||||
|
last_probe_success_at bigint,
|
||||||
|
last_probe_failure_at bigint,
|
||||||
|
probe_failure_count bigint DEFAULT 0 NOT NULL,
|
||||||
|
probe_status character varying(64) DEFAULT 'never'::character varying NOT NULL,
|
||||||
|
updated_at bigint NOT NULL
|
||||||
|
);
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
--
|
--
|
||||||
-- Name: provider_endpoints; Type: TABLE; Schema: public; Owner: -
|
-- Name: provider_endpoints; Type: TABLE; Schema: public; Owner: -
|
||||||
--
|
--
|
||||||
@@ -1659,6 +1691,21 @@ END $mig$;
|
|||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
--
|
||||||
|
-- Name: pool_member_scores pool_member_scores_pkey; Type: CONSTRAINT; Schema: public; Owner: -
|
||||||
|
--
|
||||||
|
|
||||||
|
DO $mig$ BEGIN
|
||||||
|
ALTER TABLE ONLY public.pool_member_scores
|
||||||
|
ADD CONSTRAINT pool_member_scores_pkey PRIMARY KEY (id);
|
||||||
|
EXCEPTION
|
||||||
|
WHEN duplicate_object THEN NULL;
|
||||||
|
WHEN duplicate_table THEN NULL;
|
||||||
|
WHEN invalid_table_definition THEN NULL;
|
||||||
|
END $mig$;
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
--
|
--
|
||||||
-- Name: provider_endpoints provider_endpoints_pkey; Type: CONSTRAINT; Schema: public; Owner: -
|
-- Name: provider_endpoints provider_endpoints_pkey; Type: CONSTRAINT; Schema: public; Owner: -
|
||||||
--
|
--
|
||||||
@@ -3401,6 +3448,38 @@ CREATE INDEX IF NOT EXISTS ix_provider_api_keys_id ON public.provider_api_keys U
|
|||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
--
|
||||||
|
-- Name: pool_member_scores_rank_idx; Type: INDEX; Schema: public; Owner: -
|
||||||
|
--
|
||||||
|
|
||||||
|
CREATE INDEX IF NOT EXISTS pool_member_scores_rank_idx ON public.pool_member_scores USING btree (pool_kind, pool_id, capability, scope_kind, scope_id, hard_state, score DESC);
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
--
|
||||||
|
-- Name: pool_member_scores_member_idx; Type: INDEX; Schema: public; Owner: -
|
||||||
|
--
|
||||||
|
|
||||||
|
CREATE INDEX IF NOT EXISTS pool_member_scores_member_idx ON public.pool_member_scores USING btree (pool_kind, pool_id, member_kind, member_id);
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
--
|
||||||
|
-- Name: pool_member_scores_probe_idx; Type: INDEX; Schema: public; Owner: -
|
||||||
|
--
|
||||||
|
|
||||||
|
CREATE INDEX IF NOT EXISTS pool_member_scores_probe_idx ON public.pool_member_scores USING btree (pool_kind, pool_id, probe_status, last_probe_success_at);
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
--
|
||||||
|
-- Name: pool_member_scores_updated_at_idx; Type: INDEX; Schema: public; Owner: -
|
||||||
|
--
|
||||||
|
|
||||||
|
CREATE INDEX IF NOT EXISTS pool_member_scores_updated_at_idx ON public.pool_member_scores USING btree (updated_at);
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
--
|
--
|
||||||
-- Name: ix_provider_endpoints_id; Type: INDEX; Schema: public; Owner: -
|
-- Name: ix_provider_endpoints_id; Type: INDEX; Schema: public; Owner: -
|
||||||
--
|
--
|
||||||
|
|||||||
@@ -0,0 +1,40 @@
|
|||||||
|
CREATE TABLE IF NOT EXISTS public.pool_member_scores (
|
||||||
|
id character varying(192) PRIMARY KEY,
|
||||||
|
pool_kind character varying(64) NOT NULL,
|
||||||
|
pool_id character varying(64) NOT NULL,
|
||||||
|
member_kind character varying(64) NOT NULL,
|
||||||
|
member_id character varying(64) NOT NULL,
|
||||||
|
capability character varying(64) NOT NULL,
|
||||||
|
scope_kind character varying(64) NOT NULL,
|
||||||
|
scope_id character varying(128),
|
||||||
|
score double precision NOT NULL DEFAULT 0,
|
||||||
|
hard_state character varying(64) NOT NULL DEFAULT 'unknown',
|
||||||
|
score_version bigint NOT NULL DEFAULT 1,
|
||||||
|
score_reason jsonb NOT NULL,
|
||||||
|
last_ranked_at bigint,
|
||||||
|
last_scheduled_at bigint,
|
||||||
|
last_success_at bigint,
|
||||||
|
last_failure_at bigint,
|
||||||
|
failure_count bigint NOT NULL DEFAULT 0,
|
||||||
|
last_probe_attempt_at bigint,
|
||||||
|
last_probe_success_at bigint,
|
||||||
|
last_probe_failure_at bigint,
|
||||||
|
probe_failure_count bigint NOT NULL DEFAULT 0,
|
||||||
|
probe_status character varying(64) NOT NULL DEFAULT 'never',
|
||||||
|
updated_at bigint NOT NULL
|
||||||
|
);
|
||||||
|
|
||||||
|
CREATE INDEX IF NOT EXISTS pool_member_scores_rank_idx
|
||||||
|
ON public.pool_member_scores USING btree
|
||||||
|
(pool_kind, pool_id, capability, scope_kind, scope_id, hard_state, score DESC);
|
||||||
|
|
||||||
|
CREATE INDEX IF NOT EXISTS pool_member_scores_member_idx
|
||||||
|
ON public.pool_member_scores USING btree
|
||||||
|
(pool_kind, pool_id, member_kind, member_id);
|
||||||
|
|
||||||
|
CREATE INDEX IF NOT EXISTS pool_member_scores_probe_idx
|
||||||
|
ON public.pool_member_scores USING btree
|
||||||
|
(pool_kind, pool_id, probe_status, last_probe_success_at);
|
||||||
|
|
||||||
|
CREATE INDEX IF NOT EXISTS pool_member_scores_updated_at_idx
|
||||||
|
ON public.pool_member_scores USING btree (updated_at);
|
||||||
@@ -301,6 +301,36 @@ CREATE TABLE IF NOT EXISTS provider_api_keys (
|
|||||||
);
|
);
|
||||||
CREATE INDEX IF NOT EXISTS provider_api_keys_provider_id_idx ON provider_api_keys (provider_id);
|
CREATE INDEX IF NOT EXISTS provider_api_keys_provider_id_idx ON provider_api_keys (provider_id);
|
||||||
|
|
||||||
|
CREATE TABLE IF NOT EXISTS pool_member_scores (
|
||||||
|
id TEXT PRIMARY KEY,
|
||||||
|
pool_kind TEXT NOT NULL,
|
||||||
|
pool_id TEXT NOT NULL,
|
||||||
|
member_kind TEXT NOT NULL,
|
||||||
|
member_id TEXT NOT NULL,
|
||||||
|
capability TEXT NOT NULL,
|
||||||
|
scope_kind TEXT NOT NULL,
|
||||||
|
scope_id TEXT,
|
||||||
|
score REAL NOT NULL DEFAULT 0,
|
||||||
|
hard_state TEXT NOT NULL DEFAULT 'unknown',
|
||||||
|
score_version INTEGER NOT NULL DEFAULT 1,
|
||||||
|
score_reason TEXT NOT NULL,
|
||||||
|
last_ranked_at INTEGER,
|
||||||
|
last_scheduled_at INTEGER,
|
||||||
|
last_success_at INTEGER,
|
||||||
|
last_failure_at INTEGER,
|
||||||
|
failure_count INTEGER NOT NULL DEFAULT 0,
|
||||||
|
last_probe_attempt_at INTEGER,
|
||||||
|
last_probe_success_at INTEGER,
|
||||||
|
last_probe_failure_at INTEGER,
|
||||||
|
probe_failure_count INTEGER NOT NULL DEFAULT 0,
|
||||||
|
probe_status TEXT NOT NULL DEFAULT 'never',
|
||||||
|
updated_at INTEGER NOT NULL
|
||||||
|
);
|
||||||
|
CREATE INDEX IF NOT EXISTS pool_member_scores_rank_idx ON pool_member_scores (pool_kind, pool_id, capability, scope_kind, scope_id, hard_state, score DESC);
|
||||||
|
CREATE INDEX IF NOT EXISTS pool_member_scores_member_idx ON pool_member_scores (pool_kind, pool_id, member_kind, member_id);
|
||||||
|
CREATE INDEX IF NOT EXISTS pool_member_scores_probe_idx ON pool_member_scores (pool_kind, pool_id, probe_status, last_probe_success_at);
|
||||||
|
CREATE INDEX IF NOT EXISTS pool_member_scores_updated_at_idx ON pool_member_scores (updated_at);
|
||||||
|
|
||||||
CREATE TABLE IF NOT EXISTS gemini_file_mappings (
|
CREATE TABLE IF NOT EXISTS gemini_file_mappings (
|
||||||
id TEXT PRIMARY KEY,
|
id TEXT PRIMARY KEY,
|
||||||
file_name TEXT NOT NULL UNIQUE,
|
file_name TEXT NOT NULL UNIQUE,
|
||||||
|
|||||||
@@ -0,0 +1,37 @@
|
|||||||
|
CREATE TABLE IF NOT EXISTS pool_member_scores (
|
||||||
|
id TEXT PRIMARY KEY,
|
||||||
|
pool_kind TEXT NOT NULL,
|
||||||
|
pool_id TEXT NOT NULL,
|
||||||
|
member_kind TEXT NOT NULL,
|
||||||
|
member_id TEXT NOT NULL,
|
||||||
|
capability TEXT NOT NULL,
|
||||||
|
scope_kind TEXT NOT NULL,
|
||||||
|
scope_id TEXT,
|
||||||
|
score REAL NOT NULL DEFAULT 0,
|
||||||
|
hard_state TEXT NOT NULL DEFAULT 'unknown',
|
||||||
|
score_version INTEGER NOT NULL DEFAULT 1,
|
||||||
|
score_reason TEXT NOT NULL,
|
||||||
|
last_ranked_at INTEGER,
|
||||||
|
last_scheduled_at INTEGER,
|
||||||
|
last_success_at INTEGER,
|
||||||
|
last_failure_at INTEGER,
|
||||||
|
failure_count INTEGER NOT NULL DEFAULT 0,
|
||||||
|
last_probe_attempt_at INTEGER,
|
||||||
|
last_probe_success_at INTEGER,
|
||||||
|
last_probe_failure_at INTEGER,
|
||||||
|
probe_failure_count INTEGER NOT NULL DEFAULT 0,
|
||||||
|
probe_status TEXT NOT NULL DEFAULT 'never',
|
||||||
|
updated_at INTEGER NOT NULL
|
||||||
|
);
|
||||||
|
|
||||||
|
CREATE INDEX IF NOT EXISTS pool_member_scores_rank_idx
|
||||||
|
ON pool_member_scores (pool_kind, pool_id, capability, scope_kind, scope_id, hard_state, score DESC);
|
||||||
|
|
||||||
|
CREATE INDEX IF NOT EXISTS pool_member_scores_member_idx
|
||||||
|
ON pool_member_scores (pool_kind, pool_id, member_kind, member_id);
|
||||||
|
|
||||||
|
CREATE INDEX IF NOT EXISTS pool_member_scores_probe_idx
|
||||||
|
ON pool_member_scores (pool_kind, pool_id, probe_status, last_probe_success_at);
|
||||||
|
|
||||||
|
CREATE INDEX IF NOT EXISTS pool_member_scores_updated_at_idx
|
||||||
|
ON pool_member_scores (updated_at);
|
||||||
@@ -537,6 +537,38 @@ CREATE TABLE IF NOT EXISTS public.provider_api_keys (
|
|||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
--
|
||||||
|
-- Name: pool_member_scores; Type: TABLE; Schema: public; Owner: -
|
||||||
|
--
|
||||||
|
|
||||||
|
CREATE TABLE IF NOT EXISTS public.pool_member_scores (
|
||||||
|
id character varying(192) NOT NULL,
|
||||||
|
pool_kind character varying(64) NOT NULL,
|
||||||
|
pool_id character varying(64) NOT NULL,
|
||||||
|
member_kind character varying(64) NOT NULL,
|
||||||
|
member_id character varying(64) NOT NULL,
|
||||||
|
capability character varying(64) NOT NULL,
|
||||||
|
scope_kind character varying(64) NOT NULL,
|
||||||
|
scope_id character varying(128),
|
||||||
|
score double precision DEFAULT 0 NOT NULL,
|
||||||
|
hard_state character varying(64) DEFAULT 'unknown'::character varying NOT NULL,
|
||||||
|
score_version bigint DEFAULT 1 NOT NULL,
|
||||||
|
score_reason jsonb NOT NULL,
|
||||||
|
last_ranked_at bigint,
|
||||||
|
last_scheduled_at bigint,
|
||||||
|
last_success_at bigint,
|
||||||
|
last_failure_at bigint,
|
||||||
|
failure_count bigint DEFAULT 0 NOT NULL,
|
||||||
|
last_probe_attempt_at bigint,
|
||||||
|
last_probe_success_at bigint,
|
||||||
|
last_probe_failure_at bigint,
|
||||||
|
probe_failure_count bigint DEFAULT 0 NOT NULL,
|
||||||
|
probe_status character varying(64) DEFAULT 'never'::character varying NOT NULL,
|
||||||
|
updated_at bigint NOT NULL
|
||||||
|
);
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
--
|
--
|
||||||
-- Name: provider_endpoints; Type: TABLE; Schema: public; Owner: -
|
-- Name: provider_endpoints; Type: TABLE; Schema: public; Owner: -
|
||||||
--
|
--
|
||||||
|
|||||||
@@ -252,6 +252,21 @@ END $mig$;
|
|||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
--
|
||||||
|
-- Name: pool_member_scores pool_member_scores_pkey; Type: CONSTRAINT; Schema: public; Owner: -
|
||||||
|
--
|
||||||
|
|
||||||
|
DO $mig$ BEGIN
|
||||||
|
ALTER TABLE ONLY public.pool_member_scores
|
||||||
|
ADD CONSTRAINT pool_member_scores_pkey PRIMARY KEY (id);
|
||||||
|
EXCEPTION
|
||||||
|
WHEN duplicate_object THEN NULL;
|
||||||
|
WHEN duplicate_table THEN NULL;
|
||||||
|
WHEN invalid_table_definition THEN NULL;
|
||||||
|
END $mig$;
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
--
|
--
|
||||||
-- Name: provider_endpoints provider_endpoints_pkey; Type: CONSTRAINT; Schema: public; Owner: -
|
-- Name: provider_endpoints provider_endpoints_pkey; Type: CONSTRAINT; Schema: public; Owner: -
|
||||||
--
|
--
|
||||||
|
|||||||
@@ -877,6 +877,38 @@ CREATE INDEX IF NOT EXISTS ix_provider_api_keys_id ON public.provider_api_keys U
|
|||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
--
|
||||||
|
-- Name: pool_member_scores_rank_idx; Type: INDEX; Schema: public; Owner: -
|
||||||
|
--
|
||||||
|
|
||||||
|
CREATE INDEX IF NOT EXISTS pool_member_scores_rank_idx ON public.pool_member_scores USING btree (pool_kind, pool_id, capability, scope_kind, scope_id, hard_state, score DESC);
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
--
|
||||||
|
-- Name: pool_member_scores_member_idx; Type: INDEX; Schema: public; Owner: -
|
||||||
|
--
|
||||||
|
|
||||||
|
CREATE INDEX IF NOT EXISTS pool_member_scores_member_idx ON public.pool_member_scores USING btree (pool_kind, pool_id, member_kind, member_id);
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
--
|
||||||
|
-- Name: pool_member_scores_probe_idx; Type: INDEX; Schema: public; Owner: -
|
||||||
|
--
|
||||||
|
|
||||||
|
CREATE INDEX IF NOT EXISTS pool_member_scores_probe_idx ON public.pool_member_scores USING btree (pool_kind, pool_id, probe_status, last_probe_success_at);
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
--
|
||||||
|
-- Name: pool_member_scores_updated_at_idx; Type: INDEX; Schema: public; Owner: -
|
||||||
|
--
|
||||||
|
|
||||||
|
CREATE INDEX IF NOT EXISTS pool_member_scores_updated_at_idx ON public.pool_member_scores USING btree (updated_at);
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
--
|
--
|
||||||
-- Name: ix_provider_endpoints_id; Type: INDEX; Schema: public; Owner: -
|
-- Name: ix_provider_endpoints_id; Type: INDEX; Schema: public; Owner: -
|
||||||
--
|
--
|
||||||
|
|||||||
@@ -127,6 +127,36 @@ CREATE TABLE IF NOT EXISTS provider_api_keys (
|
|||||||
KEY provider_api_keys_provider_id_idx (provider_id)
|
KEY provider_api_keys_provider_id_idx (provider_id)
|
||||||
);
|
);
|
||||||
|
|
||||||
|
CREATE TABLE IF NOT EXISTS pool_member_scores (
|
||||||
|
id VARCHAR(192) PRIMARY KEY,
|
||||||
|
pool_kind VARCHAR(64) NOT NULL,
|
||||||
|
pool_id VARCHAR(64) NOT NULL,
|
||||||
|
member_kind VARCHAR(64) NOT NULL,
|
||||||
|
member_id VARCHAR(64) NOT NULL,
|
||||||
|
capability VARCHAR(64) NOT NULL,
|
||||||
|
scope_kind VARCHAR(64) NOT NULL,
|
||||||
|
scope_id VARCHAR(128),
|
||||||
|
score DOUBLE NOT NULL DEFAULT 0,
|
||||||
|
hard_state VARCHAR(64) NOT NULL DEFAULT 'unknown',
|
||||||
|
score_version BIGINT NOT NULL DEFAULT 1,
|
||||||
|
score_reason TEXT NOT NULL,
|
||||||
|
last_ranked_at BIGINT,
|
||||||
|
last_scheduled_at BIGINT,
|
||||||
|
last_success_at BIGINT,
|
||||||
|
last_failure_at BIGINT,
|
||||||
|
failure_count BIGINT NOT NULL DEFAULT 0,
|
||||||
|
last_probe_attempt_at BIGINT,
|
||||||
|
last_probe_success_at BIGINT,
|
||||||
|
last_probe_failure_at BIGINT,
|
||||||
|
probe_failure_count BIGINT NOT NULL DEFAULT 0,
|
||||||
|
probe_status VARCHAR(64) NOT NULL DEFAULT 'never',
|
||||||
|
updated_at BIGINT NOT NULL,
|
||||||
|
KEY pool_member_scores_rank_idx (pool_kind, pool_id, capability, scope_kind, scope_id, hard_state, score DESC),
|
||||||
|
KEY pool_member_scores_member_idx (pool_kind, pool_id, member_kind, member_id),
|
||||||
|
KEY pool_member_scores_probe_idx (pool_kind, pool_id, probe_status, last_probe_success_at),
|
||||||
|
KEY pool_member_scores_updated_at_idx (updated_at)
|
||||||
|
);
|
||||||
|
|
||||||
CREATE TABLE IF NOT EXISTS gemini_file_mappings (
|
CREATE TABLE IF NOT EXISTS gemini_file_mappings (
|
||||||
id VARCHAR(64) PRIMARY KEY,
|
id VARCHAR(64) PRIMARY KEY,
|
||||||
file_name VARCHAR(512) NOT NULL,
|
file_name VARCHAR(512) NOT NULL,
|
||||||
|
|||||||
@@ -534,6 +534,38 @@ CREATE TABLE IF NOT EXISTS public.provider_api_keys (
|
|||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
--
|
||||||
|
-- Name: pool_member_scores; Type: TABLE; Schema: public; Owner: -
|
||||||
|
--
|
||||||
|
|
||||||
|
CREATE TABLE IF NOT EXISTS public.pool_member_scores (
|
||||||
|
id character varying(192) NOT NULL,
|
||||||
|
pool_kind character varying(64) NOT NULL,
|
||||||
|
pool_id character varying(64) NOT NULL,
|
||||||
|
member_kind character varying(64) NOT NULL,
|
||||||
|
member_id character varying(64) NOT NULL,
|
||||||
|
capability character varying(64) NOT NULL,
|
||||||
|
scope_kind character varying(64) NOT NULL,
|
||||||
|
scope_id character varying(128),
|
||||||
|
score double precision DEFAULT 0 NOT NULL,
|
||||||
|
hard_state character varying(64) DEFAULT 'unknown'::character varying NOT NULL,
|
||||||
|
score_version bigint DEFAULT 1 NOT NULL,
|
||||||
|
score_reason jsonb NOT NULL,
|
||||||
|
last_ranked_at bigint,
|
||||||
|
last_scheduled_at bigint,
|
||||||
|
last_success_at bigint,
|
||||||
|
last_failure_at bigint,
|
||||||
|
failure_count bigint DEFAULT 0 NOT NULL,
|
||||||
|
last_probe_attempt_at bigint,
|
||||||
|
last_probe_success_at bigint,
|
||||||
|
last_probe_failure_at bigint,
|
||||||
|
probe_failure_count bigint DEFAULT 0 NOT NULL,
|
||||||
|
probe_status character varying(64) DEFAULT 'never'::character varying NOT NULL,
|
||||||
|
updated_at bigint NOT NULL
|
||||||
|
);
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
--
|
--
|
||||||
-- Name: provider_endpoints; Type: TABLE; Schema: public; Owner: -
|
-- Name: provider_endpoints; Type: TABLE; Schema: public; Owner: -
|
||||||
--
|
--
|
||||||
|
|||||||
@@ -252,6 +252,21 @@ END $mig$;
|
|||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
--
|
||||||
|
-- Name: pool_member_scores pool_member_scores_pkey; Type: CONSTRAINT; Schema: public; Owner: -
|
||||||
|
--
|
||||||
|
|
||||||
|
DO $mig$ BEGIN
|
||||||
|
ALTER TABLE ONLY public.pool_member_scores
|
||||||
|
ADD CONSTRAINT pool_member_scores_pkey PRIMARY KEY (id);
|
||||||
|
EXCEPTION
|
||||||
|
WHEN duplicate_object THEN NULL;
|
||||||
|
WHEN duplicate_table THEN NULL;
|
||||||
|
WHEN invalid_table_definition THEN NULL;
|
||||||
|
END $mig$;
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
--
|
--
|
||||||
-- Name: provider_endpoints provider_endpoints_pkey; Type: CONSTRAINT; Schema: public; Owner: -
|
-- Name: provider_endpoints provider_endpoints_pkey; Type: CONSTRAINT; Schema: public; Owner: -
|
||||||
--
|
--
|
||||||
|
|||||||
@@ -869,6 +869,38 @@ CREATE INDEX IF NOT EXISTS ix_provider_api_keys_id ON public.provider_api_keys U
|
|||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
--
|
||||||
|
-- Name: pool_member_scores_rank_idx; Type: INDEX; Schema: public; Owner: -
|
||||||
|
--
|
||||||
|
|
||||||
|
CREATE INDEX IF NOT EXISTS pool_member_scores_rank_idx ON public.pool_member_scores USING btree (pool_kind, pool_id, capability, scope_kind, scope_id, hard_state, score DESC);
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
--
|
||||||
|
-- Name: pool_member_scores_member_idx; Type: INDEX; Schema: public; Owner: -
|
||||||
|
--
|
||||||
|
|
||||||
|
CREATE INDEX IF NOT EXISTS pool_member_scores_member_idx ON public.pool_member_scores USING btree (pool_kind, pool_id, member_kind, member_id);
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
--
|
||||||
|
-- Name: pool_member_scores_probe_idx; Type: INDEX; Schema: public; Owner: -
|
||||||
|
--
|
||||||
|
|
||||||
|
CREATE INDEX IF NOT EXISTS pool_member_scores_probe_idx ON public.pool_member_scores USING btree (pool_kind, pool_id, probe_status, last_probe_success_at);
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
--
|
||||||
|
-- Name: pool_member_scores_updated_at_idx; Type: INDEX; Schema: public; Owner: -
|
||||||
|
--
|
||||||
|
|
||||||
|
CREATE INDEX IF NOT EXISTS pool_member_scores_updated_at_idx ON public.pool_member_scores USING btree (updated_at);
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
--
|
--
|
||||||
-- Name: ix_provider_endpoints_id; Type: INDEX; Schema: public; Owner: -
|
-- Name: ix_provider_endpoints_id; Type: INDEX; Schema: public; Owner: -
|
||||||
--
|
--
|
||||||
|
|||||||
@@ -134,6 +134,36 @@ CREATE TABLE IF NOT EXISTS provider_api_keys (
|
|||||||
);
|
);
|
||||||
CREATE INDEX IF NOT EXISTS provider_api_keys_provider_id_idx ON provider_api_keys (provider_id);
|
CREATE INDEX IF NOT EXISTS provider_api_keys_provider_id_idx ON provider_api_keys (provider_id);
|
||||||
|
|
||||||
|
CREATE TABLE IF NOT EXISTS pool_member_scores (
|
||||||
|
id TEXT PRIMARY KEY,
|
||||||
|
pool_kind TEXT NOT NULL,
|
||||||
|
pool_id TEXT NOT NULL,
|
||||||
|
member_kind TEXT NOT NULL,
|
||||||
|
member_id TEXT NOT NULL,
|
||||||
|
capability TEXT NOT NULL,
|
||||||
|
scope_kind TEXT NOT NULL,
|
||||||
|
scope_id TEXT,
|
||||||
|
score REAL NOT NULL DEFAULT 0,
|
||||||
|
hard_state TEXT NOT NULL DEFAULT 'unknown',
|
||||||
|
score_version INTEGER NOT NULL DEFAULT 1,
|
||||||
|
score_reason TEXT NOT NULL,
|
||||||
|
last_ranked_at INTEGER,
|
||||||
|
last_scheduled_at INTEGER,
|
||||||
|
last_success_at INTEGER,
|
||||||
|
last_failure_at INTEGER,
|
||||||
|
failure_count INTEGER NOT NULL DEFAULT 0,
|
||||||
|
last_probe_attempt_at INTEGER,
|
||||||
|
last_probe_success_at INTEGER,
|
||||||
|
last_probe_failure_at INTEGER,
|
||||||
|
probe_failure_count INTEGER NOT NULL DEFAULT 0,
|
||||||
|
probe_status TEXT NOT NULL DEFAULT 'never',
|
||||||
|
updated_at INTEGER NOT NULL
|
||||||
|
);
|
||||||
|
CREATE INDEX IF NOT EXISTS pool_member_scores_rank_idx ON pool_member_scores (pool_kind, pool_id, capability, scope_kind, scope_id, hard_state, score DESC);
|
||||||
|
CREATE INDEX IF NOT EXISTS pool_member_scores_member_idx ON pool_member_scores (pool_kind, pool_id, member_kind, member_id);
|
||||||
|
CREATE INDEX IF NOT EXISTS pool_member_scores_probe_idx ON pool_member_scores (pool_kind, pool_id, probe_status, last_probe_success_at);
|
||||||
|
CREATE INDEX IF NOT EXISTS pool_member_scores_updated_at_idx ON pool_member_scores (updated_at);
|
||||||
|
|
||||||
CREATE TABLE IF NOT EXISTS gemini_file_mappings (
|
CREATE TABLE IF NOT EXISTS gemini_file_mappings (
|
||||||
id TEXT PRIMARY KEY,
|
id TEXT PRIMARY KEY,
|
||||||
file_name TEXT NOT NULL UNIQUE,
|
file_name TEXT NOT NULL UNIQUE,
|
||||||
|
|||||||
@@ -130,6 +130,37 @@ CREATE TABLE IF NOT EXISTS provider_api_keys (
|
|||||||
KEY provider_api_keys_provider_id_idx (`provider_id`)
|
KEY provider_api_keys_provider_id_idx (`provider_id`)
|
||||||
);
|
);
|
||||||
|
|
||||||
|
CREATE TABLE IF NOT EXISTS pool_member_scores (
|
||||||
|
`id` VARCHAR(192) NOT NULL,
|
||||||
|
`pool_kind` VARCHAR(64) NOT NULL,
|
||||||
|
`pool_id` VARCHAR(64) NOT NULL,
|
||||||
|
`member_kind` VARCHAR(64) NOT NULL,
|
||||||
|
`member_id` VARCHAR(64) NOT NULL,
|
||||||
|
`capability` VARCHAR(64) NOT NULL,
|
||||||
|
`scope_kind` VARCHAR(64) NOT NULL,
|
||||||
|
`scope_id` VARCHAR(128),
|
||||||
|
`score` DOUBLE NOT NULL DEFAULT 0,
|
||||||
|
`hard_state` VARCHAR(64) NOT NULL DEFAULT 'unknown',
|
||||||
|
`score_version` BIGINT NOT NULL DEFAULT 1,
|
||||||
|
`score_reason` JSON NOT NULL,
|
||||||
|
`last_ranked_at` BIGINT,
|
||||||
|
`last_scheduled_at` BIGINT,
|
||||||
|
`last_success_at` BIGINT,
|
||||||
|
`last_failure_at` BIGINT,
|
||||||
|
`failure_count` BIGINT NOT NULL DEFAULT 0,
|
||||||
|
`last_probe_attempt_at` BIGINT,
|
||||||
|
`last_probe_success_at` BIGINT,
|
||||||
|
`last_probe_failure_at` BIGINT,
|
||||||
|
`probe_failure_count` BIGINT NOT NULL DEFAULT 0,
|
||||||
|
`probe_status` VARCHAR(64) NOT NULL DEFAULT 'never',
|
||||||
|
`updated_at` BIGINT NOT NULL,
|
||||||
|
PRIMARY KEY (`id`),
|
||||||
|
KEY pool_member_scores_rank_idx (`pool_kind`, `pool_id`, `capability`, `scope_kind`, `scope_id`, `hard_state`, `score`),
|
||||||
|
KEY pool_member_scores_member_idx (`pool_kind`, `pool_id`, `member_kind`, `member_id`),
|
||||||
|
KEY pool_member_scores_probe_idx (`pool_kind`, `pool_id`, `probe_status`, `last_probe_success_at`),
|
||||||
|
KEY pool_member_scores_updated_at_idx (`updated_at`)
|
||||||
|
);
|
||||||
|
|
||||||
CREATE TABLE IF NOT EXISTS api_key_provider_mappings (
|
CREATE TABLE IF NOT EXISTS api_key_provider_mappings (
|
||||||
`id` VARCHAR(64) NOT NULL,
|
`id` VARCHAR(64) NOT NULL,
|
||||||
`api_key_id` VARCHAR(64) NOT NULL,
|
`api_key_id` VARCHAR(64) NOT NULL,
|
||||||
|
|||||||
@@ -134,6 +134,38 @@ CREATE TABLE IF NOT EXISTS public.provider_api_keys (
|
|||||||
ALTER TABLE ONLY public.provider_api_keys ADD CONSTRAINT provider_api_keys_pkey PRIMARY KEY (id);
|
ALTER TABLE ONLY public.provider_api_keys ADD CONSTRAINT provider_api_keys_pkey PRIMARY KEY (id);
|
||||||
CREATE INDEX IF NOT EXISTS provider_api_keys_provider_id_idx ON public.provider_api_keys USING btree (provider_id);
|
CREATE INDEX IF NOT EXISTS provider_api_keys_provider_id_idx ON public.provider_api_keys USING btree (provider_id);
|
||||||
|
|
||||||
|
CREATE TABLE IF NOT EXISTS public.pool_member_scores (
|
||||||
|
id character varying(192) NOT NULL,
|
||||||
|
pool_kind character varying(64) NOT NULL,
|
||||||
|
pool_id character varying(64) NOT NULL,
|
||||||
|
member_kind character varying(64) NOT NULL,
|
||||||
|
member_id character varying(64) NOT NULL,
|
||||||
|
capability character varying(64) NOT NULL,
|
||||||
|
scope_kind character varying(64) NOT NULL,
|
||||||
|
scope_id character varying(128),
|
||||||
|
score double precision DEFAULT 0 NOT NULL,
|
||||||
|
hard_state character varying(64) DEFAULT 'unknown' NOT NULL,
|
||||||
|
score_version bigint DEFAULT 1 NOT NULL,
|
||||||
|
score_reason jsonb NOT NULL,
|
||||||
|
last_ranked_at bigint,
|
||||||
|
last_scheduled_at bigint,
|
||||||
|
last_success_at bigint,
|
||||||
|
last_failure_at bigint,
|
||||||
|
failure_count bigint DEFAULT 0 NOT NULL,
|
||||||
|
last_probe_attempt_at bigint,
|
||||||
|
last_probe_success_at bigint,
|
||||||
|
last_probe_failure_at bigint,
|
||||||
|
probe_failure_count bigint DEFAULT 0 NOT NULL,
|
||||||
|
probe_status character varying(64) DEFAULT 'never' NOT NULL,
|
||||||
|
updated_at bigint NOT NULL
|
||||||
|
);
|
||||||
|
|
||||||
|
ALTER TABLE ONLY public.pool_member_scores ADD CONSTRAINT pool_member_scores_pkey PRIMARY KEY (id);
|
||||||
|
CREATE INDEX IF NOT EXISTS pool_member_scores_rank_idx ON public.pool_member_scores USING btree (pool_kind, pool_id, capability, scope_kind, scope_id, hard_state, score);
|
||||||
|
CREATE INDEX IF NOT EXISTS pool_member_scores_member_idx ON public.pool_member_scores USING btree (pool_kind, pool_id, member_kind, member_id);
|
||||||
|
CREATE INDEX IF NOT EXISTS pool_member_scores_probe_idx ON public.pool_member_scores USING btree (pool_kind, pool_id, probe_status, last_probe_success_at);
|
||||||
|
CREATE INDEX IF NOT EXISTS pool_member_scores_updated_at_idx ON public.pool_member_scores USING btree (updated_at);
|
||||||
|
|
||||||
CREATE TABLE IF NOT EXISTS public.api_key_provider_mappings (
|
CREATE TABLE IF NOT EXISTS public.api_key_provider_mappings (
|
||||||
id character varying(64) NOT NULL,
|
id character varying(64) NOT NULL,
|
||||||
api_key_id character varying(64) NOT NULL,
|
api_key_id character varying(64) NOT NULL,
|
||||||
|
|||||||
@@ -126,6 +126,36 @@ CREATE TABLE IF NOT EXISTS provider_api_keys (
|
|||||||
);
|
);
|
||||||
CREATE INDEX IF NOT EXISTS provider_api_keys_provider_id_idx ON provider_api_keys (provider_id);
|
CREATE INDEX IF NOT EXISTS provider_api_keys_provider_id_idx ON provider_api_keys (provider_id);
|
||||||
|
|
||||||
|
CREATE TABLE IF NOT EXISTS pool_member_scores (
|
||||||
|
id TEXT PRIMARY KEY NOT NULL,
|
||||||
|
pool_kind TEXT NOT NULL,
|
||||||
|
pool_id TEXT NOT NULL,
|
||||||
|
member_kind TEXT NOT NULL,
|
||||||
|
member_id TEXT NOT NULL,
|
||||||
|
capability TEXT NOT NULL,
|
||||||
|
scope_kind TEXT NOT NULL,
|
||||||
|
scope_id TEXT,
|
||||||
|
score REAL NOT NULL DEFAULT 0,
|
||||||
|
hard_state TEXT NOT NULL DEFAULT 'unknown',
|
||||||
|
score_version INTEGER NOT NULL DEFAULT 1,
|
||||||
|
score_reason TEXT NOT NULL,
|
||||||
|
last_ranked_at INTEGER,
|
||||||
|
last_scheduled_at INTEGER,
|
||||||
|
last_success_at INTEGER,
|
||||||
|
last_failure_at INTEGER,
|
||||||
|
failure_count INTEGER NOT NULL DEFAULT 0,
|
||||||
|
last_probe_attempt_at INTEGER,
|
||||||
|
last_probe_success_at INTEGER,
|
||||||
|
last_probe_failure_at INTEGER,
|
||||||
|
probe_failure_count INTEGER NOT NULL DEFAULT 0,
|
||||||
|
probe_status TEXT NOT NULL DEFAULT 'never',
|
||||||
|
updated_at INTEGER NOT NULL
|
||||||
|
);
|
||||||
|
CREATE INDEX IF NOT EXISTS pool_member_scores_rank_idx ON pool_member_scores (pool_kind, pool_id, capability, scope_kind, scope_id, hard_state, score);
|
||||||
|
CREATE INDEX IF NOT EXISTS pool_member_scores_member_idx ON pool_member_scores (pool_kind, pool_id, member_kind, member_id);
|
||||||
|
CREATE INDEX IF NOT EXISTS pool_member_scores_probe_idx ON pool_member_scores (pool_kind, pool_id, probe_status, last_probe_success_at);
|
||||||
|
CREATE INDEX IF NOT EXISTS pool_member_scores_updated_at_idx ON pool_member_scores (updated_at);
|
||||||
|
|
||||||
CREATE TABLE IF NOT EXISTS api_key_provider_mappings (
|
CREATE TABLE IF NOT EXISTS api_key_provider_mappings (
|
||||||
id TEXT PRIMARY KEY NOT NULL,
|
id TEXT PRIMARY KEY NOT NULL,
|
||||||
api_key_id TEXT NOT NULL,
|
api_key_id TEXT NOT NULL,
|
||||||
|
|||||||
@@ -579,6 +579,143 @@ type = "unix_seconds"
|
|||||||
name = "provider_api_keys_provider_id_idx"
|
name = "provider_api_keys_provider_id_idx"
|
||||||
columns = ["provider_id"]
|
columns = ["provider_id"]
|
||||||
|
|
||||||
|
[table.pool_member_scores]
|
||||||
|
domain = "provider_catalog"
|
||||||
|
order = 44
|
||||||
|
primary_key = ["id"]
|
||||||
|
|
||||||
|
[[table.pool_member_scores.columns]]
|
||||||
|
name = "id"
|
||||||
|
type = "text_id"
|
||||||
|
length = 192
|
||||||
|
|
||||||
|
[[table.pool_member_scores.columns]]
|
||||||
|
name = "pool_kind"
|
||||||
|
type = "text"
|
||||||
|
length = 64
|
||||||
|
|
||||||
|
[[table.pool_member_scores.columns]]
|
||||||
|
name = "pool_id"
|
||||||
|
type = "text_id"
|
||||||
|
length = 64
|
||||||
|
|
||||||
|
[[table.pool_member_scores.columns]]
|
||||||
|
name = "member_kind"
|
||||||
|
type = "text"
|
||||||
|
length = 64
|
||||||
|
|
||||||
|
[[table.pool_member_scores.columns]]
|
||||||
|
name = "member_id"
|
||||||
|
type = "text_id"
|
||||||
|
length = 64
|
||||||
|
|
||||||
|
[[table.pool_member_scores.columns]]
|
||||||
|
name = "capability"
|
||||||
|
type = "text"
|
||||||
|
length = 64
|
||||||
|
|
||||||
|
[[table.pool_member_scores.columns]]
|
||||||
|
name = "scope_kind"
|
||||||
|
type = "text"
|
||||||
|
length = 64
|
||||||
|
|
||||||
|
[[table.pool_member_scores.columns]]
|
||||||
|
name = "scope_id"
|
||||||
|
type = "text_id"
|
||||||
|
length = 128
|
||||||
|
nullable = true
|
||||||
|
|
||||||
|
[[table.pool_member_scores.columns]]
|
||||||
|
name = "score"
|
||||||
|
type = "float64"
|
||||||
|
default = 0
|
||||||
|
|
||||||
|
[[table.pool_member_scores.columns]]
|
||||||
|
name = "hard_state"
|
||||||
|
type = "text"
|
||||||
|
length = 64
|
||||||
|
default = "unknown"
|
||||||
|
|
||||||
|
[[table.pool_member_scores.columns]]
|
||||||
|
name = "score_version"
|
||||||
|
type = "int64"
|
||||||
|
default = 1
|
||||||
|
|
||||||
|
[[table.pool_member_scores.columns]]
|
||||||
|
name = "score_reason"
|
||||||
|
type = "json"
|
||||||
|
|
||||||
|
[[table.pool_member_scores.columns]]
|
||||||
|
name = "last_ranked_at"
|
||||||
|
type = "unix_seconds"
|
||||||
|
nullable = true
|
||||||
|
|
||||||
|
[[table.pool_member_scores.columns]]
|
||||||
|
name = "last_scheduled_at"
|
||||||
|
type = "unix_seconds"
|
||||||
|
nullable = true
|
||||||
|
|
||||||
|
[[table.pool_member_scores.columns]]
|
||||||
|
name = "last_success_at"
|
||||||
|
type = "unix_seconds"
|
||||||
|
nullable = true
|
||||||
|
|
||||||
|
[[table.pool_member_scores.columns]]
|
||||||
|
name = "last_failure_at"
|
||||||
|
type = "unix_seconds"
|
||||||
|
nullable = true
|
||||||
|
|
||||||
|
[[table.pool_member_scores.columns]]
|
||||||
|
name = "failure_count"
|
||||||
|
type = "int64"
|
||||||
|
default = 0
|
||||||
|
|
||||||
|
[[table.pool_member_scores.columns]]
|
||||||
|
name = "last_probe_attempt_at"
|
||||||
|
type = "unix_seconds"
|
||||||
|
nullable = true
|
||||||
|
|
||||||
|
[[table.pool_member_scores.columns]]
|
||||||
|
name = "last_probe_success_at"
|
||||||
|
type = "unix_seconds"
|
||||||
|
nullable = true
|
||||||
|
|
||||||
|
[[table.pool_member_scores.columns]]
|
||||||
|
name = "last_probe_failure_at"
|
||||||
|
type = "unix_seconds"
|
||||||
|
nullable = true
|
||||||
|
|
||||||
|
[[table.pool_member_scores.columns]]
|
||||||
|
name = "probe_failure_count"
|
||||||
|
type = "int64"
|
||||||
|
default = 0
|
||||||
|
|
||||||
|
[[table.pool_member_scores.columns]]
|
||||||
|
name = "probe_status"
|
||||||
|
type = "text"
|
||||||
|
length = 64
|
||||||
|
default = "never"
|
||||||
|
|
||||||
|
[[table.pool_member_scores.columns]]
|
||||||
|
name = "updated_at"
|
||||||
|
type = "unix_seconds"
|
||||||
|
|
||||||
|
[[table.pool_member_scores.indexes]]
|
||||||
|
name = "pool_member_scores_rank_idx"
|
||||||
|
columns = ["pool_kind", "pool_id", "capability", "scope_kind", "scope_id", "hard_state", "score"]
|
||||||
|
|
||||||
|
[[table.pool_member_scores.indexes]]
|
||||||
|
name = "pool_member_scores_member_idx"
|
||||||
|
columns = ["pool_kind", "pool_id", "member_kind", "member_id"]
|
||||||
|
|
||||||
|
[[table.pool_member_scores.indexes]]
|
||||||
|
name = "pool_member_scores_probe_idx"
|
||||||
|
columns = ["pool_kind", "pool_id", "probe_status", "last_probe_success_at"]
|
||||||
|
|
||||||
|
[[table.pool_member_scores.indexes]]
|
||||||
|
name = "pool_member_scores_updated_at_idx"
|
||||||
|
columns = ["updated_at"]
|
||||||
|
|
||||||
[table.api_key_provider_mappings]
|
[table.api_key_provider_mappings]
|
||||||
domain = "provider_catalog"
|
domain = "provider_catalog"
|
||||||
order = 45
|
order = 45
|
||||||
|
|||||||
@@ -37,6 +37,9 @@ use crate::repository::management_tokens::{
|
|||||||
use crate::repository::oauth_providers::{
|
use crate::repository::oauth_providers::{
|
||||||
MysqlOAuthProviderRepository, OAuthProviderReadRepository, OAuthProviderWriteRepository,
|
MysqlOAuthProviderRepository, OAuthProviderReadRepository, OAuthProviderWriteRepository,
|
||||||
};
|
};
|
||||||
|
use crate::repository::pool_scores::{
|
||||||
|
MysqlPoolMemberScoreRepository, PoolMemberScoreWriteRepository, PoolScoreReadRepository,
|
||||||
|
};
|
||||||
use crate::repository::provider_catalog::{
|
use crate::repository::provider_catalog::{
|
||||||
MysqlProviderCatalogReadRepository, ProviderCatalogReadRepository,
|
MysqlProviderCatalogReadRepository, ProviderCatalogReadRepository,
|
||||||
ProviderCatalogWriteRepository,
|
ProviderCatalogWriteRepository,
|
||||||
@@ -184,6 +187,14 @@ impl MysqlBackend {
|
|||||||
Arc::new(MysqlProviderCatalogReadRepository::new(self.pool_clone()))
|
Arc::new(MysqlProviderCatalogReadRepository::new(self.pool_clone()))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub fn pool_score_read_repository(&self) -> Arc<dyn PoolScoreReadRepository> {
|
||||||
|
Arc::new(MysqlPoolMemberScoreRepository::new(self.pool_clone()))
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn pool_score_write_repository(&self) -> Arc<dyn PoolMemberScoreWriteRepository> {
|
||||||
|
Arc::new(MysqlPoolMemberScoreRepository::new(self.pool_clone()))
|
||||||
|
}
|
||||||
|
|
||||||
pub fn proxy_node_read_repository(&self) -> Arc<dyn ProxyNodeReadRepository> {
|
pub fn proxy_node_read_repository(&self) -> Arc<dyn ProxyNodeReadRepository> {
|
||||||
Arc::new(MysqlProxyNodeReadRepository::new(self.pool_clone()))
|
Arc::new(MysqlProxyNodeReadRepository::new(self.pool_clone()))
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -39,6 +39,9 @@ use crate::repository::management_tokens::{
|
|||||||
use crate::repository::oauth_providers::{
|
use crate::repository::oauth_providers::{
|
||||||
OAuthProviderReadRepository, OAuthProviderWriteRepository, SqlxOAuthProviderRepository,
|
OAuthProviderReadRepository, OAuthProviderWriteRepository, SqlxOAuthProviderRepository,
|
||||||
};
|
};
|
||||||
|
use crate::repository::pool_scores::{
|
||||||
|
PoolMemberScoreWriteRepository, PoolScoreReadRepository, PostgresPoolMemberScoreRepository,
|
||||||
|
};
|
||||||
use crate::repository::provider_catalog::{
|
use crate::repository::provider_catalog::{
|
||||||
ProviderCatalogReadRepository, ProviderCatalogWriteRepository,
|
ProviderCatalogReadRepository, ProviderCatalogWriteRepository,
|
||||||
SqlxProviderCatalogReadRepository,
|
SqlxProviderCatalogReadRepository,
|
||||||
@@ -195,6 +198,14 @@ impl PostgresBackend {
|
|||||||
Arc::new(SqlxProviderCatalogReadRepository::new(self.pool_clone()))
|
Arc::new(SqlxProviderCatalogReadRepository::new(self.pool_clone()))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub fn pool_score_read_repository(&self) -> Arc<dyn PoolScoreReadRepository> {
|
||||||
|
Arc::new(PostgresPoolMemberScoreRepository::new(self.pool_clone()))
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn pool_score_write_repository(&self) -> Arc<dyn PoolMemberScoreWriteRepository> {
|
||||||
|
Arc::new(PostgresPoolMemberScoreRepository::new(self.pool_clone()))
|
||||||
|
}
|
||||||
|
|
||||||
pub fn provider_quota_read_repository(&self) -> Arc<dyn ProviderQuotaReadRepository> {
|
pub fn provider_quota_read_repository(&self) -> Arc<dyn ProviderQuotaReadRepository> {
|
||||||
Arc::new(SqlxProviderQuotaRepository::new(self.pool_clone()))
|
Arc::new(SqlxProviderQuotaRepository::new(self.pool_clone()))
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -14,6 +14,7 @@ use crate::repository::gemini_file_mappings::GeminiFileMappingReadRepository;
|
|||||||
use crate::repository::global_models::GlobalModelReadRepository;
|
use crate::repository::global_models::GlobalModelReadRepository;
|
||||||
use crate::repository::management_tokens::ManagementTokenReadRepository;
|
use crate::repository::management_tokens::ManagementTokenReadRepository;
|
||||||
use crate::repository::oauth_providers::OAuthProviderReadRepository;
|
use crate::repository::oauth_providers::OAuthProviderReadRepository;
|
||||||
|
use crate::repository::pool_scores::PoolScoreReadRepository;
|
||||||
use crate::repository::provider_catalog::ProviderCatalogReadRepository;
|
use crate::repository::provider_catalog::ProviderCatalogReadRepository;
|
||||||
use crate::repository::proxy_nodes::ProxyNodeReadRepository;
|
use crate::repository::proxy_nodes::ProxyNodeReadRepository;
|
||||||
use crate::repository::quota::ProviderQuotaReadRepository;
|
use crate::repository::quota::ProviderQuotaReadRepository;
|
||||||
@@ -34,6 +35,7 @@ pub struct DataReadRepositories {
|
|||||||
global_models: Option<Arc<dyn GlobalModelReadRepository>>,
|
global_models: Option<Arc<dyn GlobalModelReadRepository>>,
|
||||||
management_tokens: Option<Arc<dyn ManagementTokenReadRepository>>,
|
management_tokens: Option<Arc<dyn ManagementTokenReadRepository>>,
|
||||||
oauth_providers: Option<Arc<dyn OAuthProviderReadRepository>>,
|
oauth_providers: Option<Arc<dyn OAuthProviderReadRepository>>,
|
||||||
|
pool_scores: Option<Arc<dyn PoolScoreReadRepository>>,
|
||||||
proxy_nodes: Option<Arc<dyn ProxyNodeReadRepository>>,
|
proxy_nodes: Option<Arc<dyn ProxyNodeReadRepository>>,
|
||||||
minimal_candidate_selection: Option<Arc<dyn MinimalCandidateSelectionReadRepository>>,
|
minimal_candidate_selection: Option<Arc<dyn MinimalCandidateSelectionReadRepository>>,
|
||||||
request_candidates: Option<Arc<dyn RequestCandidateReadRepository>>,
|
request_candidates: Option<Arc<dyn RequestCandidateReadRepository>>,
|
||||||
@@ -61,6 +63,7 @@ impl fmt::Debug for DataReadRepositories {
|
|||||||
.field("has_global_models", &self.global_models.is_some())
|
.field("has_global_models", &self.global_models.is_some())
|
||||||
.field("has_management_tokens", &self.management_tokens.is_some())
|
.field("has_management_tokens", &self.management_tokens.is_some())
|
||||||
.field("has_oauth_providers", &self.oauth_providers.is_some())
|
.field("has_oauth_providers", &self.oauth_providers.is_some())
|
||||||
|
.field("has_pool_scores", &self.pool_scores.is_some())
|
||||||
.field("has_proxy_nodes", &self.proxy_nodes.is_some())
|
.field("has_proxy_nodes", &self.proxy_nodes.is_some())
|
||||||
.field(
|
.field(
|
||||||
"has_minimal_candidate_selection",
|
"has_minimal_candidate_selection",
|
||||||
@@ -124,6 +127,10 @@ impl DataReadRepositories {
|
|||||||
.map(PostgresBackend::oauth_provider_read_repository)
|
.map(PostgresBackend::oauth_provider_read_repository)
|
||||||
.or_else(|| mysql.map(MysqlBackend::oauth_provider_read_repository))
|
.or_else(|| mysql.map(MysqlBackend::oauth_provider_read_repository))
|
||||||
.or_else(|| sqlite.map(SqliteBackend::oauth_provider_read_repository)),
|
.or_else(|| sqlite.map(SqliteBackend::oauth_provider_read_repository)),
|
||||||
|
pool_scores: postgres
|
||||||
|
.map(PostgresBackend::pool_score_read_repository)
|
||||||
|
.or_else(|| mysql.map(MysqlBackend::pool_score_read_repository))
|
||||||
|
.or_else(|| sqlite.map(SqliteBackend::pool_score_read_repository)),
|
||||||
proxy_nodes: postgres
|
proxy_nodes: postgres
|
||||||
.map(PostgresBackend::proxy_node_read_repository)
|
.map(PostgresBackend::proxy_node_read_repository)
|
||||||
.or_else(|| mysql.map(MysqlBackend::proxy_node_read_repository))
|
.or_else(|| mysql.map(MysqlBackend::proxy_node_read_repository))
|
||||||
@@ -208,6 +215,10 @@ impl DataReadRepositories {
|
|||||||
self.oauth_providers.clone()
|
self.oauth_providers.clone()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub fn pool_scores(&self) -> Option<Arc<dyn PoolScoreReadRepository>> {
|
||||||
|
self.pool_scores.clone()
|
||||||
|
}
|
||||||
|
|
||||||
pub fn proxy_nodes(&self) -> Option<Arc<dyn ProxyNodeReadRepository>> {
|
pub fn proxy_nodes(&self) -> Option<Arc<dyn ProxyNodeReadRepository>> {
|
||||||
self.proxy_nodes.clone()
|
self.proxy_nodes.clone()
|
||||||
}
|
}
|
||||||
@@ -257,6 +268,7 @@ impl DataReadRepositories {
|
|||||||
|| self.global_models.is_some()
|
|| self.global_models.is_some()
|
||||||
|| self.management_tokens.is_some()
|
|| self.management_tokens.is_some()
|
||||||
|| self.oauth_providers.is_some()
|
|| self.oauth_providers.is_some()
|
||||||
|
|| self.pool_scores.is_some()
|
||||||
|| self.proxy_nodes.is_some()
|
|| self.proxy_nodes.is_some()
|
||||||
|| self.minimal_candidate_selection.is_some()
|
|| self.minimal_candidate_selection.is_some()
|
||||||
|| self.request_candidates.is_some()
|
|| self.request_candidates.is_some()
|
||||||
|
|||||||
@@ -37,6 +37,9 @@ use crate::repository::management_tokens::{
|
|||||||
use crate::repository::oauth_providers::{
|
use crate::repository::oauth_providers::{
|
||||||
OAuthProviderReadRepository, OAuthProviderWriteRepository, SqliteOAuthProviderRepository,
|
OAuthProviderReadRepository, OAuthProviderWriteRepository, SqliteOAuthProviderRepository,
|
||||||
};
|
};
|
||||||
|
use crate::repository::pool_scores::{
|
||||||
|
PoolMemberScoreWriteRepository, PoolScoreReadRepository, SqlitePoolMemberScoreRepository,
|
||||||
|
};
|
||||||
use crate::repository::provider_catalog::{
|
use crate::repository::provider_catalog::{
|
||||||
ProviderCatalogReadRepository, ProviderCatalogWriteRepository,
|
ProviderCatalogReadRepository, ProviderCatalogWriteRepository,
|
||||||
SqliteProviderCatalogReadRepository,
|
SqliteProviderCatalogReadRepository,
|
||||||
@@ -197,6 +200,14 @@ impl SqliteBackend {
|
|||||||
Arc::new(SqliteProviderCatalogReadRepository::new(self.pool_clone()))
|
Arc::new(SqliteProviderCatalogReadRepository::new(self.pool_clone()))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub fn pool_score_read_repository(&self) -> Arc<dyn PoolScoreReadRepository> {
|
||||||
|
Arc::new(SqlitePoolMemberScoreRepository::new(self.pool_clone()))
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn pool_score_write_repository(&self) -> Arc<dyn PoolMemberScoreWriteRepository> {
|
||||||
|
Arc::new(SqlitePoolMemberScoreRepository::new(self.pool_clone()))
|
||||||
|
}
|
||||||
|
|
||||||
pub fn proxy_node_read_repository(&self) -> Arc<dyn ProxyNodeReadRepository> {
|
pub fn proxy_node_read_repository(&self) -> Arc<dyn ProxyNodeReadRepository> {
|
||||||
Arc::new(SqliteProxyNodeReadRepository::new(self.pool_clone()))
|
Arc::new(SqliteProxyNodeReadRepository::new(self.pool_clone()))
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -11,6 +11,7 @@ use crate::repository::gemini_file_mappings::GeminiFileMappingWriteRepository;
|
|||||||
use crate::repository::global_models::GlobalModelWriteRepository;
|
use crate::repository::global_models::GlobalModelWriteRepository;
|
||||||
use crate::repository::management_tokens::ManagementTokenWriteRepository;
|
use crate::repository::management_tokens::ManagementTokenWriteRepository;
|
||||||
use crate::repository::oauth_providers::OAuthProviderWriteRepository;
|
use crate::repository::oauth_providers::OAuthProviderWriteRepository;
|
||||||
|
use crate::repository::pool_scores::PoolMemberScoreWriteRepository;
|
||||||
use crate::repository::provider_catalog::ProviderCatalogWriteRepository;
|
use crate::repository::provider_catalog::ProviderCatalogWriteRepository;
|
||||||
use crate::repository::proxy_nodes::ProxyNodeWriteRepository;
|
use crate::repository::proxy_nodes::ProxyNodeWriteRepository;
|
||||||
use crate::repository::quota::ProviderQuotaWriteRepository;
|
use crate::repository::quota::ProviderQuotaWriteRepository;
|
||||||
@@ -30,6 +31,7 @@ pub struct DataWriteRepositories {
|
|||||||
global_models: Option<Arc<dyn GlobalModelWriteRepository>>,
|
global_models: Option<Arc<dyn GlobalModelWriteRepository>>,
|
||||||
management_tokens: Option<Arc<dyn ManagementTokenWriteRepository>>,
|
management_tokens: Option<Arc<dyn ManagementTokenWriteRepository>>,
|
||||||
oauth_providers: Option<Arc<dyn OAuthProviderWriteRepository>>,
|
oauth_providers: Option<Arc<dyn OAuthProviderWriteRepository>>,
|
||||||
|
pool_scores: Option<Arc<dyn PoolMemberScoreWriteRepository>>,
|
||||||
proxy_nodes: Option<Arc<dyn ProxyNodeWriteRepository>>,
|
proxy_nodes: Option<Arc<dyn ProxyNodeWriteRepository>>,
|
||||||
provider_catalog: Option<Arc<dyn ProviderCatalogWriteRepository>>,
|
provider_catalog: Option<Arc<dyn ProviderCatalogWriteRepository>>,
|
||||||
provider_quotas: Option<Arc<dyn ProviderQuotaWriteRepository>>,
|
provider_quotas: Option<Arc<dyn ProviderQuotaWriteRepository>>,
|
||||||
@@ -54,6 +56,7 @@ impl fmt::Debug for DataWriteRepositories {
|
|||||||
.field("has_global_models", &self.global_models.is_some())
|
.field("has_global_models", &self.global_models.is_some())
|
||||||
.field("has_management_tokens", &self.management_tokens.is_some())
|
.field("has_management_tokens", &self.management_tokens.is_some())
|
||||||
.field("has_oauth_providers", &self.oauth_providers.is_some())
|
.field("has_oauth_providers", &self.oauth_providers.is_some())
|
||||||
|
.field("has_pool_scores", &self.pool_scores.is_some())
|
||||||
.field("has_proxy_nodes", &self.proxy_nodes.is_some())
|
.field("has_proxy_nodes", &self.proxy_nodes.is_some())
|
||||||
.field("has_provider_catalog", &self.provider_catalog.is_some())
|
.field("has_provider_catalog", &self.provider_catalog.is_some())
|
||||||
.field("has_provider_quotas", &self.provider_quotas.is_some())
|
.field("has_provider_quotas", &self.provider_quotas.is_some())
|
||||||
@@ -108,6 +111,10 @@ impl DataWriteRepositories {
|
|||||||
.map(PostgresBackend::oauth_provider_write_repository)
|
.map(PostgresBackend::oauth_provider_write_repository)
|
||||||
.or_else(|| mysql.map(MysqlBackend::oauth_provider_write_repository))
|
.or_else(|| mysql.map(MysqlBackend::oauth_provider_write_repository))
|
||||||
.or_else(|| sqlite.map(SqliteBackend::oauth_provider_write_repository)),
|
.or_else(|| sqlite.map(SqliteBackend::oauth_provider_write_repository)),
|
||||||
|
pool_scores: postgres
|
||||||
|
.map(PostgresBackend::pool_score_write_repository)
|
||||||
|
.or_else(|| mysql.map(MysqlBackend::pool_score_write_repository))
|
||||||
|
.or_else(|| sqlite.map(SqliteBackend::pool_score_write_repository)),
|
||||||
proxy_nodes: postgres
|
proxy_nodes: postgres
|
||||||
.map(PostgresBackend::proxy_node_write_repository)
|
.map(PostgresBackend::proxy_node_write_repository)
|
||||||
.or_else(|| mysql.map(MysqlBackend::proxy_node_write_repository))
|
.or_else(|| mysql.map(MysqlBackend::proxy_node_write_repository))
|
||||||
@@ -184,6 +191,10 @@ impl DataWriteRepositories {
|
|||||||
self.oauth_providers.clone()
|
self.oauth_providers.clone()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub fn pool_scores(&self) -> Option<Arc<dyn PoolMemberScoreWriteRepository>> {
|
||||||
|
self.pool_scores.clone()
|
||||||
|
}
|
||||||
|
|
||||||
pub fn proxy_nodes(&self) -> Option<Arc<dyn ProxyNodeWriteRepository>> {
|
pub fn proxy_nodes(&self) -> Option<Arc<dyn ProxyNodeWriteRepository>> {
|
||||||
self.proxy_nodes.clone()
|
self.proxy_nodes.clone()
|
||||||
}
|
}
|
||||||
@@ -218,6 +229,7 @@ impl DataWriteRepositories {
|
|||||||
|| self.global_models.is_some()
|
|| self.global_models.is_some()
|
||||||
|| self.management_tokens.is_some()
|
|| self.management_tokens.is_some()
|
||||||
|| self.oauth_providers.is_some()
|
|| self.oauth_providers.is_some()
|
||||||
|
|| self.pool_scores.is_some()
|
||||||
|| self.proxy_nodes.is_some()
|
|| self.proxy_nodes.is_some()
|
||||||
|| self.provider_catalog.is_some()
|
|| self.provider_catalog.is_some()
|
||||||
|| self.provider_quotas.is_some()
|
|| self.provider_quotas.is_some()
|
||||||
|
|||||||
@@ -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 = 20260511120000;
|
pub(crate) const EMPTY_DATABASE_SNAPSHOT_CUTOFF_VERSION: i64 = 20260511130000;
|
||||||
|
|
||||||
const PUBLIC_BASE_TABLE_COUNT_SQL: &str = r#"
|
const PUBLIC_BASE_TABLE_COUNT_SQL: &str = r#"
|
||||||
SELECT COUNT(*)::BIGINT
|
SELECT COUNT(*)::BIGINT
|
||||||
|
|||||||
@@ -300,6 +300,7 @@ fn empty_database_snapshot_covers_current_cutoff_versions() {
|
|||||||
20260510120000,
|
20260510120000,
|
||||||
20260511000000,
|
20260511000000,
|
||||||
20260511120000,
|
20260511120000,
|
||||||
|
20260511130000,
|
||||||
]
|
]
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
@@ -564,7 +565,8 @@ fn mysql_and_sqlite_migrations_include_enabled_incrementals() {
|
|||||||
20260509000000,
|
20260509000000,
|
||||||
20260509120000,
|
20260509120000,
|
||||||
20260510120000,
|
20260510120000,
|
||||||
20260511120000
|
20260511120000,
|
||||||
|
20260511130000,
|
||||||
]
|
]
|
||||||
);
|
);
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
@@ -576,7 +578,8 @@ fn mysql_and_sqlite_migrations_include_enabled_incrementals() {
|
|||||||
20260509000000,
|
20260509000000,
|
||||||
20260509120000,
|
20260509120000,
|
||||||
20260510120000,
|
20260510120000,
|
||||||
20260511120000
|
20260511120000,
|
||||||
|
20260511130000,
|
||||||
]
|
]
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
@@ -1088,6 +1091,7 @@ fn pending_migrations_from_applied_skips_versions_already_applied() {
|
|||||||
20260510120000,
|
20260510120000,
|
||||||
20260511000000,
|
20260511000000,
|
||||||
20260511120000,
|
20260511120000,
|
||||||
|
20260511130000,
|
||||||
]
|
]
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -4,8 +4,8 @@ use async_trait::async_trait;
|
|||||||
|
|
||||||
use super::{
|
use super::{
|
||||||
MinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
|
MinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
|
||||||
StoredPoolKeyCandidateOrder, StoredPoolKeyCandidateRowsQuery,
|
StoredPoolKeyCandidateOrder, StoredPoolKeyCandidateRowsByKeyIdsQuery,
|
||||||
StoredRequestedModelCandidateRowsQuery,
|
StoredPoolKeyCandidateRowsQuery, StoredRequestedModelCandidateRowsQuery,
|
||||||
};
|
};
|
||||||
use crate::DataLayerError;
|
use crate::DataLayerError;
|
||||||
|
|
||||||
@@ -138,6 +138,39 @@ impl MinimalCandidateSelectionReadRepository for InMemoryMinimalCandidateSelecti
|
|||||||
.take(query.limit as usize)
|
.take(query.limit as usize)
|
||||||
.collect())
|
.collect())
|
||||||
}
|
}
|
||||||
|
|
||||||
|
async fn list_pool_key_rows_for_group_key_ids(
|
||||||
|
&self,
|
||||||
|
query: &StoredPoolKeyCandidateRowsByKeyIdsQuery,
|
||||||
|
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
|
||||||
|
if query.key_ids.is_empty() {
|
||||||
|
return Ok(Vec::new());
|
||||||
|
}
|
||||||
|
let key_order = query
|
||||||
|
.key_ids
|
||||||
|
.iter()
|
||||||
|
.enumerate()
|
||||||
|
.map(|(index, key_id)| (key_id.as_str(), index))
|
||||||
|
.collect::<std::collections::BTreeMap<_, _>>();
|
||||||
|
let mut rows = self
|
||||||
|
.list_for_exact_api_format(&query.api_format)
|
||||||
|
.await?
|
||||||
|
.into_iter()
|
||||||
|
.filter(|row| {
|
||||||
|
row.provider_id == query.provider_id
|
||||||
|
&& row.endpoint_id == query.endpoint_id
|
||||||
|
&& row.model_id == query.model_id
|
||||||
|
&& key_order.contains_key(row.key_id.as_str())
|
||||||
|
})
|
||||||
|
.collect::<Vec<_>>();
|
||||||
|
rows.sort_by(|left, right| {
|
||||||
|
key_order
|
||||||
|
.get(left.key_id.as_str())
|
||||||
|
.cmp(&key_order.get(right.key_id.as_str()))
|
||||||
|
.then(left.key_id.cmp(&right.key_id))
|
||||||
|
});
|
||||||
|
Ok(rows)
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
fn sort_pool_key_rows(
|
fn sort_pool_key_rows(
|
||||||
|
|||||||
@@ -7,8 +7,8 @@ mod sqlite;
|
|||||||
pub(crate) use aether_data_contracts::repository::candidate_selection::{
|
pub(crate) use aether_data_contracts::repository::candidate_selection::{
|
||||||
MinimalCandidateSelectionReadRepository, MinimalCandidateSelectionRepository,
|
MinimalCandidateSelectionReadRepository, MinimalCandidateSelectionRepository,
|
||||||
StoredMinimalCandidateSelectionRow, StoredPoolKeyCandidateOrder,
|
StoredMinimalCandidateSelectionRow, StoredPoolKeyCandidateOrder,
|
||||||
StoredPoolKeyCandidateRowsQuery, StoredProviderModelMapping,
|
StoredPoolKeyCandidateRowsByKeyIdsQuery, StoredPoolKeyCandidateRowsQuery,
|
||||||
StoredRequestedModelCandidateRowsQuery,
|
StoredProviderModelMapping, StoredRequestedModelCandidateRowsQuery,
|
||||||
};
|
};
|
||||||
pub use memory::InMemoryMinimalCandidateSelectionReadRepository;
|
pub use memory::InMemoryMinimalCandidateSelectionReadRepository;
|
||||||
pub use mysql::MysqlMinimalCandidateSelectionReadRepository;
|
pub use mysql::MysqlMinimalCandidateSelectionReadRepository;
|
||||||
|
|||||||
@@ -5,7 +5,8 @@ use sqlx::{mysql::MySqlRow, MySql, QueryBuilder, Row};
|
|||||||
|
|
||||||
use super::{
|
use super::{
|
||||||
MinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
|
MinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
|
||||||
StoredPoolKeyCandidateOrder, StoredPoolKeyCandidateRowsQuery, StoredProviderModelMapping,
|
StoredPoolKeyCandidateOrder, StoredPoolKeyCandidateRowsByKeyIdsQuery,
|
||||||
|
StoredPoolKeyCandidateRowsQuery, StoredProviderModelMapping,
|
||||||
StoredRequestedModelCandidateRowsQuery,
|
StoredRequestedModelCandidateRowsQuery,
|
||||||
};
|
};
|
||||||
use crate::driver::mysql::MysqlPool;
|
use crate::driver::mysql::MysqlPool;
|
||||||
@@ -196,6 +197,40 @@ impl MinimalCandidateSelectionReadRepository for MysqlMinimalCandidateSelectionR
|
|||||||
.map(|item| item.row)
|
.map(|item| item.row)
|
||||||
.collect())
|
.collect())
|
||||||
}
|
}
|
||||||
|
|
||||||
|
async fn list_pool_key_rows_for_group_key_ids(
|
||||||
|
&self,
|
||||||
|
query: &StoredPoolKeyCandidateRowsByKeyIdsQuery,
|
||||||
|
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
|
||||||
|
if query.key_ids.is_empty() {
|
||||||
|
return Ok(Vec::new());
|
||||||
|
}
|
||||||
|
let key_order = query
|
||||||
|
.key_ids
|
||||||
|
.iter()
|
||||||
|
.enumerate()
|
||||||
|
.map(|(index, key_id)| (key_id.as_str(), index))
|
||||||
|
.collect::<BTreeMap<_, _>>();
|
||||||
|
let mut rows = self
|
||||||
|
.load_rows_for_api_format(&query.api_format)
|
||||||
|
.await?
|
||||||
|
.into_iter()
|
||||||
|
.filter(|row| {
|
||||||
|
row.row.provider_id == query.provider_id
|
||||||
|
&& row.row.endpoint_id == query.endpoint_id
|
||||||
|
&& row.row.model_id == query.model_id
|
||||||
|
&& key_order.contains_key(row.row.key_id.as_str())
|
||||||
|
})
|
||||||
|
.map(|item| item.row)
|
||||||
|
.collect::<Vec<_>>();
|
||||||
|
rows.sort_by(|left, right| {
|
||||||
|
key_order
|
||||||
|
.get(left.key_id.as_str())
|
||||||
|
.cmp(&key_order.get(right.key_id.as_str()))
|
||||||
|
.then(left.key_id.cmp(&right.key_id))
|
||||||
|
});
|
||||||
|
Ok(dedupe_candidate_selection_rows(rows))
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
fn select_pool_rows(rows: Vec<CandidateSelectionRow>) -> Vec<StoredMinimalCandidateSelectionRow> {
|
fn select_pool_rows(rows: Vec<CandidateSelectionRow>) -> Vec<StoredMinimalCandidateSelectionRow> {
|
||||||
|
|||||||
@@ -5,7 +5,8 @@ use std::collections::BTreeSet;
|
|||||||
|
|
||||||
use super::{
|
use super::{
|
||||||
MinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
|
MinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
|
||||||
StoredPoolKeyCandidateOrder, StoredPoolKeyCandidateRowsQuery, StoredProviderModelMapping,
|
StoredPoolKeyCandidateOrder, StoredPoolKeyCandidateRowsByKeyIdsQuery,
|
||||||
|
StoredPoolKeyCandidateRowsQuery, StoredProviderModelMapping,
|
||||||
StoredRequestedModelCandidateRowsQuery,
|
StoredRequestedModelCandidateRowsQuery,
|
||||||
};
|
};
|
||||||
use crate::{error::SqlxResultExt, DataLayerError};
|
use crate::{error::SqlxResultExt, DataLayerError};
|
||||||
@@ -535,6 +536,13 @@ fn pool_key_candidate_selection_sql(order: &StoredPoolKeyCandidateOrder) -> Stri
|
|||||||
LIST_POOL_KEYS_FOR_GROUP_SQL.replace(default_order, &replacement)
|
LIST_POOL_KEYS_FOR_GROUP_SQL.replace(default_order, &replacement)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn pool_key_candidate_selection_by_key_ids_sql() -> String {
|
||||||
|
let default_order =
|
||||||
|
"ORDER BY\n pak.internal_priority ASC,\n pak.id ASC\nLIMIT $7\nOFFSET $8\n";
|
||||||
|
let replacement = "AND pak.id = ANY($7::text[])\nORDER BY\n array_position($7::text[], pak.id) ASC,\n pak.id ASC\n";
|
||||||
|
LIST_POOL_KEYS_FOR_GROUP_SQL.replace(default_order, replacement)
|
||||||
|
}
|
||||||
|
|
||||||
#[derive(Debug, Clone)]
|
#[derive(Debug, Clone)]
|
||||||
pub struct SqlxMinimalCandidateSelectionReadRepository {
|
pub struct SqlxMinimalCandidateSelectionReadRepository {
|
||||||
pool: PgPool,
|
pool: PgPool,
|
||||||
@@ -704,6 +712,51 @@ impl SqlxMinimalCandidateSelectionReadRepository {
|
|||||||
}
|
}
|
||||||
Ok(dedupe_candidate_selection_rows(rows))
|
Ok(dedupe_candidate_selection_rows(rows))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub async fn list_pool_key_rows_for_group_key_ids(
|
||||||
|
&self,
|
||||||
|
query: &StoredPoolKeyCandidateRowsByKeyIdsQuery,
|
||||||
|
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
|
||||||
|
if query.key_ids.is_empty() {
|
||||||
|
return Ok(Vec::new());
|
||||||
|
}
|
||||||
|
let mut rows = Vec::new();
|
||||||
|
let canonical_api_format = normalize_api_format(&query.api_format);
|
||||||
|
let storage_aliases = api_format_aliases(&canonical_api_format);
|
||||||
|
let sql_match_aliases = sql_match_aliases(&storage_aliases);
|
||||||
|
let sql = pool_key_candidate_selection_by_key_ids_sql();
|
||||||
|
for api_format in storage_aliases {
|
||||||
|
rows.extend(
|
||||||
|
Self::collect_query_rows(
|
||||||
|
sqlx::query(sql.as_str())
|
||||||
|
.bind(api_format)
|
||||||
|
.bind(query.provider_id.as_str())
|
||||||
|
.bind(query.endpoint_id.as_str())
|
||||||
|
.bind(query.model_id.as_str())
|
||||||
|
.bind(sql_match_aliases.clone())
|
||||||
|
.bind(canonical_api_format.clone())
|
||||||
|
.bind(query.key_ids.clone())
|
||||||
|
.fetch(&self.pool),
|
||||||
|
map_candidate_selection_row,
|
||||||
|
)
|
||||||
|
.await?,
|
||||||
|
);
|
||||||
|
}
|
||||||
|
let key_order = query
|
||||||
|
.key_ids
|
||||||
|
.iter()
|
||||||
|
.enumerate()
|
||||||
|
.map(|(index, key_id)| (key_id.as_str(), index))
|
||||||
|
.collect::<std::collections::BTreeMap<_, _>>();
|
||||||
|
let mut rows = dedupe_candidate_selection_rows(rows);
|
||||||
|
rows.sort_by(|left, right| {
|
||||||
|
key_order
|
||||||
|
.get(left.key_id.as_str())
|
||||||
|
.cmp(&key_order.get(right.key_id.as_str()))
|
||||||
|
.then(left.key_id.cmp(&right.key_id))
|
||||||
|
});
|
||||||
|
Ok(rows)
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
fn requested_model_selection_sql() -> String {
|
fn requested_model_selection_sql() -> String {
|
||||||
@@ -917,6 +970,13 @@ impl MinimalCandidateSelectionReadRepository for SqlxMinimalCandidateSelectionRe
|
|||||||
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
|
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
|
||||||
Self::list_pool_key_rows_for_group(self, query).await
|
Self::list_pool_key_rows_for_group(self, query).await
|
||||||
}
|
}
|
||||||
|
|
||||||
|
async fn list_pool_key_rows_for_group_key_ids(
|
||||||
|
&self,
|
||||||
|
query: &StoredPoolKeyCandidateRowsByKeyIdsQuery,
|
||||||
|
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
|
||||||
|
Self::list_pool_key_rows_for_group_key_ids(self, query).await
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
fn map_candidate_selection_row(
|
fn map_candidate_selection_row(
|
||||||
|
|||||||
@@ -5,7 +5,8 @@ use sqlx::{sqlite::SqliteRow, QueryBuilder, Row, Sqlite};
|
|||||||
|
|
||||||
use super::{
|
use super::{
|
||||||
MinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
|
MinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
|
||||||
StoredPoolKeyCandidateOrder, StoredPoolKeyCandidateRowsQuery, StoredProviderModelMapping,
|
StoredPoolKeyCandidateOrder, StoredPoolKeyCandidateRowsByKeyIdsQuery,
|
||||||
|
StoredPoolKeyCandidateRowsQuery, StoredProviderModelMapping,
|
||||||
StoredRequestedModelCandidateRowsQuery,
|
StoredRequestedModelCandidateRowsQuery,
|
||||||
};
|
};
|
||||||
use crate::driver::sqlite::SqlitePool;
|
use crate::driver::sqlite::SqlitePool;
|
||||||
@@ -196,6 +197,40 @@ impl MinimalCandidateSelectionReadRepository for SqliteMinimalCandidateSelection
|
|||||||
.map(|item| item.row)
|
.map(|item| item.row)
|
||||||
.collect())
|
.collect())
|
||||||
}
|
}
|
||||||
|
|
||||||
|
async fn list_pool_key_rows_for_group_key_ids(
|
||||||
|
&self,
|
||||||
|
query: &StoredPoolKeyCandidateRowsByKeyIdsQuery,
|
||||||
|
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
|
||||||
|
if query.key_ids.is_empty() {
|
||||||
|
return Ok(Vec::new());
|
||||||
|
}
|
||||||
|
let key_order = query
|
||||||
|
.key_ids
|
||||||
|
.iter()
|
||||||
|
.enumerate()
|
||||||
|
.map(|(index, key_id)| (key_id.as_str(), index))
|
||||||
|
.collect::<BTreeMap<_, _>>();
|
||||||
|
let mut rows = self
|
||||||
|
.load_rows_for_api_format(&query.api_format)
|
||||||
|
.await?
|
||||||
|
.into_iter()
|
||||||
|
.filter(|row| {
|
||||||
|
row.row.provider_id == query.provider_id
|
||||||
|
&& row.row.endpoint_id == query.endpoint_id
|
||||||
|
&& row.row.model_id == query.model_id
|
||||||
|
&& key_order.contains_key(row.row.key_id.as_str())
|
||||||
|
})
|
||||||
|
.map(|item| item.row)
|
||||||
|
.collect::<Vec<_>>();
|
||||||
|
rows.sort_by(|left, right| {
|
||||||
|
key_order
|
||||||
|
.get(left.key_id.as_str())
|
||||||
|
.cmp(&key_order.get(right.key_id.as_str()))
|
||||||
|
.then(left.key_id.cmp(&right.key_id))
|
||||||
|
});
|
||||||
|
Ok(dedupe_candidate_selection_rows(rows))
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
fn select_pool_rows(rows: Vec<CandidateSelectionRow>) -> Vec<StoredMinimalCandidateSelectionRow> {
|
fn select_pool_rows(rows: Vec<CandidateSelectionRow>) -> Vec<StoredMinimalCandidateSelectionRow> {
|
||||||
|
|||||||
@@ -16,6 +16,7 @@ pub mod gemini_file_mappings;
|
|||||||
pub mod global_models;
|
pub mod global_models;
|
||||||
pub mod management_tokens;
|
pub mod management_tokens;
|
||||||
pub mod oauth_providers;
|
pub mod oauth_providers;
|
||||||
|
pub mod pool_scores;
|
||||||
pub mod provider_catalog;
|
pub mod provider_catalog;
|
||||||
pub mod provider_oauth;
|
pub mod provider_oauth;
|
||||||
pub mod proxy_nodes;
|
pub mod proxy_nodes;
|
||||||
|
|||||||
513
crates/aether-data/src/repository/pool_scores/memory.rs
Normal file
513
crates/aether-data/src/repository/pool_scores/memory.rs
Normal file
@@ -0,0 +1,513 @@
|
|||||||
|
use std::collections::BTreeMap;
|
||||||
|
use std::sync::RwLock;
|
||||||
|
|
||||||
|
use async_trait::async_trait;
|
||||||
|
|
||||||
|
use super::{
|
||||||
|
score_with_delta, GetPoolMemberScoresByIdsQuery, ListPoolMemberProbeCandidatesQuery,
|
||||||
|
ListPoolMemberScoresQuery, ListRankedPoolMembersQuery, PoolMemberHardState, PoolMemberIdentity,
|
||||||
|
PoolMemberProbeAttempt, PoolMemberProbeResult, PoolMemberProbeStatus,
|
||||||
|
PoolMemberScheduleFeedback, PoolMemberScoreWriteRepository, PoolScoreReadRepository,
|
||||||
|
PoolScoreScope, StoredPoolMemberScore, UpsertPoolMemberScore,
|
||||||
|
};
|
||||||
|
use crate::repository::pool_scores::merge_score_reason_patch;
|
||||||
|
use crate::DataLayerError;
|
||||||
|
|
||||||
|
#[derive(Debug, Default)]
|
||||||
|
pub struct InMemoryPoolMemberScoreRepository {
|
||||||
|
scores: RwLock<BTreeMap<String, StoredPoolMemberScore>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl InMemoryPoolMemberScoreRepository {
|
||||||
|
pub fn seed<I>(scores: I) -> Self
|
||||||
|
where
|
||||||
|
I: IntoIterator<Item = StoredPoolMemberScore>,
|
||||||
|
{
|
||||||
|
Self {
|
||||||
|
scores: RwLock::new(
|
||||||
|
scores
|
||||||
|
.into_iter()
|
||||||
|
.map(|score| (score.id.clone(), score))
|
||||||
|
.collect(),
|
||||||
|
),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn matches_identity(score: &StoredPoolMemberScore, identity: &PoolMemberIdentity) -> bool {
|
||||||
|
score.pool_kind == identity.pool_kind
|
||||||
|
&& score.pool_id == identity.pool_id
|
||||||
|
&& score.member_kind == identity.member_kind
|
||||||
|
&& score.member_id == identity.member_id
|
||||||
|
}
|
||||||
|
|
||||||
|
fn matches_scope(score: &StoredPoolMemberScore, scope: Option<&PoolScoreScope>) -> bool {
|
||||||
|
let Some(scope) = scope else {
|
||||||
|
return true;
|
||||||
|
};
|
||||||
|
score.capability == scope.capability
|
||||||
|
&& score.scope_kind == scope.scope_kind
|
||||||
|
&& score.scope_id == scope.scope_id
|
||||||
|
}
|
||||||
|
|
||||||
|
fn sort_ranked(scores: &mut [StoredPoolMemberScore]) {
|
||||||
|
scores.sort_by(|left, right| {
|
||||||
|
right
|
||||||
|
.score
|
||||||
|
.partial_cmp(&left.score)
|
||||||
|
.unwrap_or(std::cmp::Ordering::Equal)
|
||||||
|
.then_with(|| {
|
||||||
|
right
|
||||||
|
.last_ranked_at
|
||||||
|
.unwrap_or(0)
|
||||||
|
.cmp(&left.last_ranked_at.unwrap_or(0))
|
||||||
|
})
|
||||||
|
.then_with(|| left.member_id.cmp(&right.member_id))
|
||||||
|
.then_with(|| left.id.cmp(&right.id))
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
fn sort_probe(scores: &mut [StoredPoolMemberScore]) {
|
||||||
|
scores.sort_by(|left, right| {
|
||||||
|
probe_priority(left)
|
||||||
|
.cmp(&probe_priority(right))
|
||||||
|
.then_with(|| right.probe_failure_count.cmp(&left.probe_failure_count))
|
||||||
|
.then_with(|| {
|
||||||
|
left.last_probe_success_at
|
||||||
|
.unwrap_or(0)
|
||||||
|
.cmp(&right.last_probe_success_at.unwrap_or(0))
|
||||||
|
})
|
||||||
|
.then_with(|| {
|
||||||
|
left.last_scheduled_at
|
||||||
|
.unwrap_or(0)
|
||||||
|
.cmp(&right.last_scheduled_at.unwrap_or(0))
|
||||||
|
.reverse()
|
||||||
|
})
|
||||||
|
.then_with(|| left.member_id.cmp(&right.member_id))
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl PoolScoreReadRepository for InMemoryPoolMemberScoreRepository {
|
||||||
|
async fn list_ranked_pool_members(
|
||||||
|
&self,
|
||||||
|
query: &ListRankedPoolMembersQuery,
|
||||||
|
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
|
||||||
|
let hard_states = query
|
||||||
|
.hard_states
|
||||||
|
.iter()
|
||||||
|
.copied()
|
||||||
|
.collect::<std::collections::BTreeSet<_>>();
|
||||||
|
let probe_statuses = query.probe_statuses.as_ref().map(|items| {
|
||||||
|
items
|
||||||
|
.iter()
|
||||||
|
.copied()
|
||||||
|
.collect::<std::collections::BTreeSet<_>>()
|
||||||
|
});
|
||||||
|
let mut scores = self
|
||||||
|
.scores
|
||||||
|
.read()
|
||||||
|
.expect("pool member score repository lock")
|
||||||
|
.values()
|
||||||
|
.filter(|score| {
|
||||||
|
score.pool_kind == query.pool_kind
|
||||||
|
&& score.pool_id == query.pool_id
|
||||||
|
&& score.capability == query.capability
|
||||||
|
&& score.scope_kind == query.scope_kind
|
||||||
|
&& score.scope_id == query.scope_id
|
||||||
|
&& (hard_states.is_empty() || hard_states.contains(&score.hard_state))
|
||||||
|
&& probe_statuses
|
||||||
|
.as_ref()
|
||||||
|
.is_none_or(|statuses| statuses.contains(&score.probe_status))
|
||||||
|
})
|
||||||
|
.cloned()
|
||||||
|
.collect::<Vec<_>>();
|
||||||
|
Self::sort_ranked(&mut scores);
|
||||||
|
Ok(scores
|
||||||
|
.into_iter()
|
||||||
|
.skip(query.offset)
|
||||||
|
.take(query.limit.max(1))
|
||||||
|
.collect())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn list_pool_member_scores(
|
||||||
|
&self,
|
||||||
|
query: &ListPoolMemberScoresQuery,
|
||||||
|
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
|
||||||
|
let hard_states = query
|
||||||
|
.hard_states
|
||||||
|
.iter()
|
||||||
|
.copied()
|
||||||
|
.collect::<std::collections::BTreeSet<_>>();
|
||||||
|
let probe_statuses = query.probe_statuses.as_ref().map(|items| {
|
||||||
|
items
|
||||||
|
.iter()
|
||||||
|
.copied()
|
||||||
|
.collect::<std::collections::BTreeSet<_>>()
|
||||||
|
});
|
||||||
|
let mut scores = self
|
||||||
|
.scores
|
||||||
|
.read()
|
||||||
|
.expect("pool member score repository lock")
|
||||||
|
.values()
|
||||||
|
.filter(|score| {
|
||||||
|
score.pool_kind == query.pool_kind
|
||||||
|
&& score.pool_id == query.pool_id
|
||||||
|
&& query
|
||||||
|
.capability
|
||||||
|
.as_ref()
|
||||||
|
.is_none_or(|capability| score.capability == *capability)
|
||||||
|
&& query
|
||||||
|
.scope_kind
|
||||||
|
.as_ref()
|
||||||
|
.is_none_or(|scope_kind| score.scope_kind == *scope_kind)
|
||||||
|
&& query
|
||||||
|
.scope_id
|
||||||
|
.as_ref()
|
||||||
|
.is_none_or(|scope_id| score.scope_id.as_ref() == Some(scope_id))
|
||||||
|
&& (hard_states.is_empty() || hard_states.contains(&score.hard_state))
|
||||||
|
&& probe_statuses
|
||||||
|
.as_ref()
|
||||||
|
.is_none_or(|statuses| statuses.contains(&score.probe_status))
|
||||||
|
})
|
||||||
|
.cloned()
|
||||||
|
.collect::<Vec<_>>();
|
||||||
|
Self::sort_ranked(&mut scores);
|
||||||
|
Ok(scores
|
||||||
|
.into_iter()
|
||||||
|
.skip(query.offset)
|
||||||
|
.take(query.limit.max(1))
|
||||||
|
.collect())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn list_pool_member_probe_candidates(
|
||||||
|
&self,
|
||||||
|
query: &ListPoolMemberProbeCandidatesQuery,
|
||||||
|
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
|
||||||
|
let mut scores = self
|
||||||
|
.scores
|
||||||
|
.read()
|
||||||
|
.expect("pool member score repository lock")
|
||||||
|
.values()
|
||||||
|
.filter(|score| {
|
||||||
|
score.pool_kind == query.pool_kind
|
||||||
|
&& score.pool_id == query.pool_id
|
||||||
|
&& query
|
||||||
|
.capability
|
||||||
|
.as_ref()
|
||||||
|
.is_none_or(|capability| score.capability == *capability)
|
||||||
|
&& matches!(
|
||||||
|
score.hard_state,
|
||||||
|
PoolMemberHardState::Available
|
||||||
|
| PoolMemberHardState::Unknown
|
||||||
|
| PoolMemberHardState::Cooldown
|
||||||
|
| PoolMemberHardState::QuotaExhausted
|
||||||
|
)
|
||||||
|
&& match score.probe_status {
|
||||||
|
PoolMemberProbeStatus::Never
|
||||||
|
| PoolMemberProbeStatus::Failed
|
||||||
|
| PoolMemberProbeStatus::Stale => true,
|
||||||
|
PoolMemberProbeStatus::Ok => score
|
||||||
|
.last_probe_success_at
|
||||||
|
.is_none_or(|ts| ts <= query.stale_before_unix_secs),
|
||||||
|
PoolMemberProbeStatus::InProgress => score
|
||||||
|
.last_probe_attempt_at
|
||||||
|
.is_none_or(|ts| ts <= query.stale_before_unix_secs),
|
||||||
|
}
|
||||||
|
})
|
||||||
|
.cloned()
|
||||||
|
.collect::<Vec<_>>();
|
||||||
|
Self::sort_probe(&mut scores);
|
||||||
|
Ok(scores.into_iter().take(query.limit.max(1)).collect())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn get_pool_member_scores_by_ids(
|
||||||
|
&self,
|
||||||
|
query: &GetPoolMemberScoresByIdsQuery,
|
||||||
|
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
|
||||||
|
let ids = query
|
||||||
|
.ids
|
||||||
|
.iter()
|
||||||
|
.cloned()
|
||||||
|
.collect::<std::collections::BTreeSet<_>>();
|
||||||
|
let scores = self
|
||||||
|
.scores
|
||||||
|
.read()
|
||||||
|
.expect("pool member score repository lock")
|
||||||
|
.values()
|
||||||
|
.filter(|score| ids.contains(&score.id))
|
||||||
|
.cloned()
|
||||||
|
.collect::<Vec<_>>();
|
||||||
|
Ok(scores)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl PoolMemberScoreWriteRepository for InMemoryPoolMemberScoreRepository {
|
||||||
|
async fn upsert_pool_member_score(
|
||||||
|
&self,
|
||||||
|
score: UpsertPoolMemberScore,
|
||||||
|
) -> Result<StoredPoolMemberScore, DataLayerError> {
|
||||||
|
score.validate()?;
|
||||||
|
let stored = score.into_stored();
|
||||||
|
self.scores
|
||||||
|
.write()
|
||||||
|
.expect("pool member score repository lock")
|
||||||
|
.insert(stored.id.clone(), stored.clone());
|
||||||
|
Ok(stored)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn mark_pool_member_probe_in_progress(
|
||||||
|
&self,
|
||||||
|
attempt: PoolMemberProbeAttempt,
|
||||||
|
) -> Result<usize, DataLayerError> {
|
||||||
|
let mut updated = 0;
|
||||||
|
let mut guard = self
|
||||||
|
.scores
|
||||||
|
.write()
|
||||||
|
.expect("pool member score repository lock");
|
||||||
|
for score in guard.values_mut() {
|
||||||
|
if !Self::matches_identity(score, &attempt.identity)
|
||||||
|
|| !Self::matches_scope(score, attempt.scope.as_ref())
|
||||||
|
{
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
score.last_probe_attempt_at = Some(attempt.attempted_at);
|
||||||
|
score.probe_status = PoolMemberProbeStatus::InProgress;
|
||||||
|
score.score_reason = merge_score_reason_patch(
|
||||||
|
score.score_reason.clone(),
|
||||||
|
attempt.score_reason_patch.clone(),
|
||||||
|
);
|
||||||
|
score.updated_at = attempt.attempted_at;
|
||||||
|
updated += 1;
|
||||||
|
}
|
||||||
|
Ok(updated)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn record_pool_member_probe_result(
|
||||||
|
&self,
|
||||||
|
result: PoolMemberProbeResult,
|
||||||
|
) -> Result<usize, DataLayerError> {
|
||||||
|
let mut updated = 0;
|
||||||
|
let mut guard = self
|
||||||
|
.scores
|
||||||
|
.write()
|
||||||
|
.expect("pool member score repository lock");
|
||||||
|
for score in guard.values_mut() {
|
||||||
|
if !Self::matches_identity(score, &result.identity)
|
||||||
|
|| !Self::matches_scope(score, result.scope.as_ref())
|
||||||
|
{
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
score.last_probe_attempt_at = Some(result.attempted_at);
|
||||||
|
score.probe_status = result.probe_status;
|
||||||
|
if result.succeeded {
|
||||||
|
score.last_probe_success_at = Some(result.attempted_at);
|
||||||
|
score.probe_failure_count = 0;
|
||||||
|
} else {
|
||||||
|
score.last_probe_failure_at = Some(result.attempted_at);
|
||||||
|
score.probe_failure_count = score.probe_failure_count.saturating_add(1);
|
||||||
|
}
|
||||||
|
if let Some(hard_state) = result.hard_state {
|
||||||
|
score.hard_state = hard_state;
|
||||||
|
}
|
||||||
|
score.score_reason = merge_score_reason_patch(
|
||||||
|
score.score_reason.clone(),
|
||||||
|
result.score_reason_patch.clone(),
|
||||||
|
);
|
||||||
|
score.updated_at = result.attempted_at;
|
||||||
|
updated += 1;
|
||||||
|
}
|
||||||
|
Ok(updated)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn record_pool_member_schedule_feedback(
|
||||||
|
&self,
|
||||||
|
feedback: PoolMemberScheduleFeedback,
|
||||||
|
) -> Result<usize, DataLayerError> {
|
||||||
|
let mut updated = 0;
|
||||||
|
let mut guard = self
|
||||||
|
.scores
|
||||||
|
.write()
|
||||||
|
.expect("pool member score repository lock");
|
||||||
|
for score in guard.values_mut() {
|
||||||
|
if !Self::matches_identity(score, &feedback.identity)
|
||||||
|
|| !Self::matches_scope(score, feedback.scope.as_ref())
|
||||||
|
{
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
score.last_scheduled_at = Some(feedback.scheduled_at);
|
||||||
|
match feedback.succeeded {
|
||||||
|
Some(true) => {
|
||||||
|
score.last_success_at = Some(feedback.scheduled_at);
|
||||||
|
}
|
||||||
|
Some(false) => {
|
||||||
|
score.last_failure_at = Some(feedback.scheduled_at);
|
||||||
|
score.failure_count = score.failure_count.saturating_add(1);
|
||||||
|
}
|
||||||
|
None => {}
|
||||||
|
}
|
||||||
|
if let Some(hard_state) = feedback.hard_state {
|
||||||
|
score.hard_state = hard_state;
|
||||||
|
}
|
||||||
|
score.score = score_with_delta(score.score, feedback.score_delta);
|
||||||
|
score.score_reason = merge_score_reason_patch(
|
||||||
|
score.score_reason.clone(),
|
||||||
|
feedback.score_reason_patch.clone(),
|
||||||
|
);
|
||||||
|
score.updated_at = feedback.scheduled_at;
|
||||||
|
updated += 1;
|
||||||
|
}
|
||||||
|
Ok(updated)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn mark_pool_member_hard_state(
|
||||||
|
&self,
|
||||||
|
identity: &PoolMemberIdentity,
|
||||||
|
scope: Option<&PoolScoreScope>,
|
||||||
|
hard_state: PoolMemberHardState,
|
||||||
|
updated_at: u64,
|
||||||
|
) -> Result<usize, DataLayerError> {
|
||||||
|
let mut updated = 0;
|
||||||
|
let mut guard = self
|
||||||
|
.scores
|
||||||
|
.write()
|
||||||
|
.expect("pool member score repository lock");
|
||||||
|
for score in guard.values_mut() {
|
||||||
|
if Self::matches_identity(score, identity) && Self::matches_scope(score, scope) {
|
||||||
|
score.hard_state = hard_state;
|
||||||
|
score.updated_at = updated_at;
|
||||||
|
updated += 1;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(updated)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn delete_pool_member_scores_for_member(
|
||||||
|
&self,
|
||||||
|
identity: &PoolMemberIdentity,
|
||||||
|
) -> Result<usize, DataLayerError> {
|
||||||
|
let mut guard = self
|
||||||
|
.scores
|
||||||
|
.write()
|
||||||
|
.expect("pool member score repository lock");
|
||||||
|
let before = guard.len();
|
||||||
|
guard.retain(|_, score| !Self::matches_identity(score, identity));
|
||||||
|
Ok(before.saturating_sub(guard.len()))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn probe_priority(score: &StoredPoolMemberScore) -> u8 {
|
||||||
|
if score.last_scheduled_at.is_some() && score.probe_status != PoolMemberProbeStatus::Ok {
|
||||||
|
return 0;
|
||||||
|
}
|
||||||
|
match score.hard_state {
|
||||||
|
PoolMemberHardState::QuotaExhausted => 1,
|
||||||
|
PoolMemberHardState::Unknown => 2,
|
||||||
|
_ if score.probe_status == PoolMemberProbeStatus::Stale => 3,
|
||||||
|
_ => 4,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
use crate::repository::pool_scores::{
|
||||||
|
POOL_KIND_PROVIDER_KEY_POOL, POOL_MEMBER_KIND_PROVIDER_API_KEY,
|
||||||
|
POOL_SCORE_CAPABILITY_API_FORMAT, POOL_SCORE_SCOPE_KIND_MODEL,
|
||||||
|
};
|
||||||
|
|
||||||
|
fn score(id: &str, member_id: &str, value: f64) -> StoredPoolMemberScore {
|
||||||
|
StoredPoolMemberScore {
|
||||||
|
id: id.to_string(),
|
||||||
|
pool_kind: POOL_KIND_PROVIDER_KEY_POOL.to_string(),
|
||||||
|
pool_id: "provider-1".to_string(),
|
||||||
|
member_kind: POOL_MEMBER_KIND_PROVIDER_API_KEY.to_string(),
|
||||||
|
member_id: member_id.to_string(),
|
||||||
|
capability: POOL_SCORE_CAPABILITY_API_FORMAT.to_string(),
|
||||||
|
scope_kind: POOL_SCORE_SCOPE_KIND_MODEL.to_string(),
|
||||||
|
scope_id: Some("model-1".to_string()),
|
||||||
|
score: value,
|
||||||
|
hard_state: PoolMemberHardState::Available,
|
||||||
|
score_version: 1,
|
||||||
|
score_reason: serde_json::json!({}),
|
||||||
|
last_ranked_at: Some(1),
|
||||||
|
last_scheduled_at: None,
|
||||||
|
last_success_at: None,
|
||||||
|
last_failure_at: None,
|
||||||
|
failure_count: 0,
|
||||||
|
last_probe_attempt_at: None,
|
||||||
|
last_probe_success_at: None,
|
||||||
|
last_probe_failure_at: None,
|
||||||
|
probe_failure_count: 0,
|
||||||
|
probe_status: PoolMemberProbeStatus::Never,
|
||||||
|
updated_at: 1,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn lists_ranked_members_by_score() {
|
||||||
|
let repository = InMemoryPoolMemberScoreRepository::seed(vec![
|
||||||
|
score("score-1", "key-1", 0.2),
|
||||||
|
score("score-2", "key-2", 0.9),
|
||||||
|
]);
|
||||||
|
|
||||||
|
let rows = repository
|
||||||
|
.list_ranked_pool_members(&ListRankedPoolMembersQuery {
|
||||||
|
pool_kind: POOL_KIND_PROVIDER_KEY_POOL.to_string(),
|
||||||
|
pool_id: "provider-1".to_string(),
|
||||||
|
capability: POOL_SCORE_CAPABILITY_API_FORMAT.to_string(),
|
||||||
|
scope_kind: POOL_SCORE_SCOPE_KIND_MODEL.to_string(),
|
||||||
|
scope_id: Some("model-1".to_string()),
|
||||||
|
hard_states: vec![PoolMemberHardState::Available],
|
||||||
|
probe_statuses: None,
|
||||||
|
offset: 0,
|
||||||
|
limit: 10,
|
||||||
|
})
|
||||||
|
.await
|
||||||
|
.expect("list should succeed");
|
||||||
|
|
||||||
|
assert_eq!(
|
||||||
|
rows.into_iter()
|
||||||
|
.map(|row| row.member_id)
|
||||||
|
.collect::<Vec<_>>(),
|
||||||
|
vec!["key-2".to_string(), "key-1".to_string()]
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn marks_probe_in_progress_without_incrementing_failure_count() {
|
||||||
|
let repository =
|
||||||
|
InMemoryPoolMemberScoreRepository::seed(vec![score("score-1", "key-1", 0.2)]);
|
||||||
|
|
||||||
|
let updated = repository
|
||||||
|
.mark_pool_member_probe_in_progress(PoolMemberProbeAttempt {
|
||||||
|
identity: PoolMemberIdentity::provider_api_key("provider-1", "key-1"),
|
||||||
|
scope: None,
|
||||||
|
attempted_at: 100,
|
||||||
|
score_reason_patch: Some(serde_json::json!({ "last_probe": "in_progress" })),
|
||||||
|
})
|
||||||
|
.await
|
||||||
|
.expect("mark should succeed");
|
||||||
|
|
||||||
|
assert_eq!(updated, 1);
|
||||||
|
let rows = repository
|
||||||
|
.list_pool_member_scores(&ListPoolMemberScoresQuery {
|
||||||
|
pool_kind: POOL_KIND_PROVIDER_KEY_POOL.to_string(),
|
||||||
|
pool_id: "provider-1".to_string(),
|
||||||
|
capability: None,
|
||||||
|
scope_kind: None,
|
||||||
|
scope_id: None,
|
||||||
|
hard_states: Vec::new(),
|
||||||
|
probe_statuses: Some(vec![PoolMemberProbeStatus::InProgress]),
|
||||||
|
offset: 0,
|
||||||
|
limit: 10,
|
||||||
|
})
|
||||||
|
.await
|
||||||
|
.expect("list should succeed");
|
||||||
|
|
||||||
|
assert_eq!(rows.len(), 1);
|
||||||
|
assert_eq!(rows[0].last_probe_attempt_at, Some(100));
|
||||||
|
assert_eq!(rows[0].probe_failure_count, 0);
|
||||||
|
}
|
||||||
|
}
|
||||||
54
crates/aether-data/src/repository/pool_scores/mod.rs
Normal file
54
crates/aether-data/src/repository/pool_scores/mod.rs
Normal file
@@ -0,0 +1,54 @@
|
|||||||
|
pub use aether_data_contracts::repository::pool_scores::*;
|
||||||
|
|
||||||
|
mod memory;
|
||||||
|
mod mysql;
|
||||||
|
mod postgres;
|
||||||
|
mod sqlite;
|
||||||
|
|
||||||
|
pub use memory::InMemoryPoolMemberScoreRepository;
|
||||||
|
pub use mysql::MysqlPoolMemberScoreRepository;
|
||||||
|
pub use postgres::PostgresPoolMemberScoreRepository;
|
||||||
|
pub use sqlite::SqlitePoolMemberScoreRepository;
|
||||||
|
|
||||||
|
fn merge_score_reason_patch(
|
||||||
|
mut current: serde_json::Value,
|
||||||
|
patch: Option<serde_json::Value>,
|
||||||
|
) -> serde_json::Value {
|
||||||
|
let Some(patch) = patch else {
|
||||||
|
return current;
|
||||||
|
};
|
||||||
|
match (current.as_object_mut(), patch) {
|
||||||
|
(Some(current), serde_json::Value::Object(patch)) => {
|
||||||
|
for (key, value) in patch {
|
||||||
|
current.insert(key, value);
|
||||||
|
}
|
||||||
|
serde_json::Value::Object(current.clone())
|
||||||
|
}
|
||||||
|
(_, patch) => patch,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn score_with_delta(score: f64, delta_basis_points: Option<i32>) -> f64 {
|
||||||
|
let delta = delta_basis_points.unwrap_or_default() as f64 / 10_000.0;
|
||||||
|
(score + delta).clamp(0.0, 1.0)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn i64_from_u64(value: u64, field: &str) -> Result<i64, crate::DataLayerError> {
|
||||||
|
i64::try_from(value).map_err(|_| {
|
||||||
|
crate::DataLayerError::InvalidInput(format!("{field} exceeds signed 64-bit range"))
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn i64_opt_from_u64(value: Option<u64>, field: &str) -> Result<Option<i64>, crate::DataLayerError> {
|
||||||
|
value.map(|value| i64_from_u64(value, field)).transpose()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn u64_from_i64(value: i64, field: &str) -> Result<u64, crate::DataLayerError> {
|
||||||
|
u64::try_from(value).map_err(|_| {
|
||||||
|
crate::DataLayerError::UnexpectedValue(format!("{field} is negative: {value}"))
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn u64_opt_from_i64(value: Option<i64>, field: &str) -> Result<Option<u64>, crate::DataLayerError> {
|
||||||
|
value.map(|value| u64_from_i64(value, field)).transpose()
|
||||||
|
}
|
||||||
586
crates/aether-data/src/repository/pool_scores/mysql.rs
Normal file
586
crates/aether-data/src/repository/pool_scores/mysql.rs
Normal file
@@ -0,0 +1,586 @@
|
|||||||
|
use async_trait::async_trait;
|
||||||
|
use sqlx::{mysql::MySqlRow, MySql, MySqlPool, QueryBuilder, Row};
|
||||||
|
|
||||||
|
use super::{
|
||||||
|
i64_from_u64, i64_opt_from_u64, score_with_delta, u64_from_i64, u64_opt_from_i64,
|
||||||
|
GetPoolMemberScoresByIdsQuery, ListPoolMemberProbeCandidatesQuery, ListPoolMemberScoresQuery,
|
||||||
|
ListRankedPoolMembersQuery, PoolMemberHardState, PoolMemberIdentity, PoolMemberProbeAttempt,
|
||||||
|
PoolMemberProbeResult, PoolMemberProbeStatus, PoolMemberScheduleFeedback,
|
||||||
|
PoolMemberScoreWriteRepository, PoolScoreReadRepository, PoolScoreScope, StoredPoolMemberScore,
|
||||||
|
UpsertPoolMemberScore,
|
||||||
|
};
|
||||||
|
use crate::error::SqlResultExt;
|
||||||
|
use crate::repository::pool_scores::merge_score_reason_patch;
|
||||||
|
use crate::DataLayerError;
|
||||||
|
|
||||||
|
const SCORE_COLUMNS: &str = r#"
|
||||||
|
SELECT
|
||||||
|
id,
|
||||||
|
pool_kind,
|
||||||
|
pool_id,
|
||||||
|
member_kind,
|
||||||
|
member_id,
|
||||||
|
capability,
|
||||||
|
scope_kind,
|
||||||
|
scope_id,
|
||||||
|
score,
|
||||||
|
hard_state,
|
||||||
|
score_version,
|
||||||
|
score_reason,
|
||||||
|
last_ranked_at,
|
||||||
|
last_scheduled_at,
|
||||||
|
last_success_at,
|
||||||
|
last_failure_at,
|
||||||
|
failure_count,
|
||||||
|
last_probe_attempt_at,
|
||||||
|
last_probe_success_at,
|
||||||
|
last_probe_failure_at,
|
||||||
|
probe_failure_count,
|
||||||
|
probe_status,
|
||||||
|
updated_at
|
||||||
|
FROM pool_member_scores
|
||||||
|
"#;
|
||||||
|
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct MysqlPoolMemberScoreRepository {
|
||||||
|
pool: MySqlPool,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl MysqlPoolMemberScoreRepository {
|
||||||
|
pub fn new(pool: MySqlPool) -> Self {
|
||||||
|
Self { pool }
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn find_scores_by_identity(
|
||||||
|
&self,
|
||||||
|
identity: &PoolMemberIdentity,
|
||||||
|
scope: Option<&PoolScoreScope>,
|
||||||
|
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
|
||||||
|
let mut builder = QueryBuilder::<MySql>::new(SCORE_COLUMNS);
|
||||||
|
builder
|
||||||
|
.push(" WHERE pool_kind = ")
|
||||||
|
.push_bind(identity.pool_kind.clone())
|
||||||
|
.push(" AND pool_id = ")
|
||||||
|
.push_bind(identity.pool_id.clone())
|
||||||
|
.push(" AND member_kind = ")
|
||||||
|
.push_bind(identity.member_kind.clone())
|
||||||
|
.push(" AND member_id = ")
|
||||||
|
.push_bind(identity.member_id.clone());
|
||||||
|
if let Some(scope) = scope {
|
||||||
|
builder
|
||||||
|
.push(" AND capability = ")
|
||||||
|
.push_bind(scope.capability.clone())
|
||||||
|
.push(" AND scope_kind = ")
|
||||||
|
.push_bind(scope.scope_kind.clone());
|
||||||
|
if let Some(scope_id) = &scope.scope_id {
|
||||||
|
builder.push(" AND scope_id = ").push_bind(scope_id.clone());
|
||||||
|
} else {
|
||||||
|
builder.push(" AND scope_id IS NULL");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?;
|
||||||
|
rows.iter().map(map_score_row).collect()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl PoolScoreReadRepository for MysqlPoolMemberScoreRepository {
|
||||||
|
async fn list_ranked_pool_members(
|
||||||
|
&self,
|
||||||
|
query: &ListRankedPoolMembersQuery,
|
||||||
|
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
|
||||||
|
let mut builder = QueryBuilder::<MySql>::new(SCORE_COLUMNS);
|
||||||
|
builder
|
||||||
|
.push(" WHERE pool_kind = ")
|
||||||
|
.push_bind(query.pool_kind.clone())
|
||||||
|
.push(" AND pool_id = ")
|
||||||
|
.push_bind(query.pool_id.clone())
|
||||||
|
.push(" AND capability = ")
|
||||||
|
.push_bind(query.capability.clone())
|
||||||
|
.push(" AND scope_kind = ")
|
||||||
|
.push_bind(query.scope_kind.clone());
|
||||||
|
if let Some(scope_id) = &query.scope_id {
|
||||||
|
builder.push(" AND scope_id = ").push_bind(scope_id.clone());
|
||||||
|
} else {
|
||||||
|
builder.push(" AND scope_id IS NULL");
|
||||||
|
}
|
||||||
|
if !query.hard_states.is_empty() {
|
||||||
|
builder.push(" AND hard_state IN (");
|
||||||
|
let mut separated = builder.separated(", ");
|
||||||
|
for state in &query.hard_states {
|
||||||
|
separated.push_bind(state.as_database());
|
||||||
|
}
|
||||||
|
separated.push_unseparated(")");
|
||||||
|
}
|
||||||
|
if let Some(statuses) = &query.probe_statuses {
|
||||||
|
if !statuses.is_empty() {
|
||||||
|
builder.push(" AND probe_status IN (");
|
||||||
|
let mut separated = builder.separated(", ");
|
||||||
|
for status in statuses {
|
||||||
|
separated.push_bind(status.as_database());
|
||||||
|
}
|
||||||
|
separated.push_unseparated(")");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
builder
|
||||||
|
.push(" ORDER BY score DESC, last_ranked_at DESC, member_id ASC, id ASC")
|
||||||
|
.push(" LIMIT ")
|
||||||
|
.push_bind(i64_from_usize(query.limit.max(1), "pool score limit")?)
|
||||||
|
.push(" OFFSET ")
|
||||||
|
.push_bind(i64_from_usize(query.offset, "pool score offset")?);
|
||||||
|
let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?;
|
||||||
|
rows.iter().map(map_score_row).collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn list_pool_member_scores(
|
||||||
|
&self,
|
||||||
|
query: &ListPoolMemberScoresQuery,
|
||||||
|
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
|
||||||
|
let mut builder = QueryBuilder::<MySql>::new(SCORE_COLUMNS);
|
||||||
|
builder
|
||||||
|
.push(" WHERE pool_kind = ")
|
||||||
|
.push_bind(query.pool_kind.clone())
|
||||||
|
.push(" AND pool_id = ")
|
||||||
|
.push_bind(query.pool_id.clone());
|
||||||
|
if let Some(capability) = &query.capability {
|
||||||
|
builder
|
||||||
|
.push(" AND capability = ")
|
||||||
|
.push_bind(capability.clone());
|
||||||
|
}
|
||||||
|
if let Some(scope_kind) = &query.scope_kind {
|
||||||
|
builder
|
||||||
|
.push(" AND scope_kind = ")
|
||||||
|
.push_bind(scope_kind.clone());
|
||||||
|
}
|
||||||
|
if let Some(scope_id) = &query.scope_id {
|
||||||
|
builder.push(" AND scope_id = ").push_bind(scope_id.clone());
|
||||||
|
}
|
||||||
|
if !query.hard_states.is_empty() {
|
||||||
|
builder.push(" AND hard_state IN (");
|
||||||
|
let mut separated = builder.separated(", ");
|
||||||
|
for state in &query.hard_states {
|
||||||
|
separated.push_bind(state.as_database());
|
||||||
|
}
|
||||||
|
separated.push_unseparated(")");
|
||||||
|
}
|
||||||
|
if let Some(statuses) = &query.probe_statuses {
|
||||||
|
if !statuses.is_empty() {
|
||||||
|
builder.push(" AND probe_status IN (");
|
||||||
|
let mut separated = builder.separated(", ");
|
||||||
|
for status in statuses {
|
||||||
|
separated.push_bind(status.as_database());
|
||||||
|
}
|
||||||
|
separated.push_unseparated(")");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
builder
|
||||||
|
.push(" ORDER BY score DESC, last_ranked_at DESC, member_id ASC, id ASC")
|
||||||
|
.push(" LIMIT ")
|
||||||
|
.push_bind(i64_from_usize(query.limit.max(1), "pool score limit")?)
|
||||||
|
.push(" OFFSET ")
|
||||||
|
.push_bind(i64_from_usize(query.offset, "pool score offset")?);
|
||||||
|
let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?;
|
||||||
|
rows.iter().map(map_score_row).collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn list_pool_member_probe_candidates(
|
||||||
|
&self,
|
||||||
|
query: &ListPoolMemberProbeCandidatesQuery,
|
||||||
|
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
|
||||||
|
let mut builder = QueryBuilder::<MySql>::new(SCORE_COLUMNS);
|
||||||
|
builder
|
||||||
|
.push(" WHERE pool_kind = ")
|
||||||
|
.push_bind(query.pool_kind.clone())
|
||||||
|
.push(" AND pool_id = ")
|
||||||
|
.push_bind(query.pool_id.clone());
|
||||||
|
if let Some(capability) = &query.capability {
|
||||||
|
builder
|
||||||
|
.push(" AND capability = ")
|
||||||
|
.push_bind(capability.clone());
|
||||||
|
}
|
||||||
|
builder
|
||||||
|
.push(" AND hard_state IN ('available','unknown','cooldown','quota_exhausted')")
|
||||||
|
.push(" AND (probe_status IN ('never','failed','stale')")
|
||||||
|
.push(" OR (probe_status = 'ok' AND (last_probe_success_at IS NULL OR last_probe_success_at <= ")
|
||||||
|
.push_bind(i64_from_u64(
|
||||||
|
query.stale_before_unix_secs,
|
||||||
|
"pool probe stale_before_unix_secs",
|
||||||
|
)?)
|
||||||
|
.push("))")
|
||||||
|
.push(" OR (probe_status = 'in_progress' AND (last_probe_attempt_at IS NULL OR last_probe_attempt_at <= ")
|
||||||
|
.push_bind(i64_from_u64(
|
||||||
|
query.stale_before_unix_secs,
|
||||||
|
"pool probe stale_before_unix_secs",
|
||||||
|
)?)
|
||||||
|
.push(")))")
|
||||||
|
.push(
|
||||||
|
r#"
|
||||||
|
ORDER BY
|
||||||
|
CASE
|
||||||
|
WHEN last_scheduled_at IS NOT NULL AND probe_status <> 'ok' THEN 0
|
||||||
|
WHEN hard_state = 'quota_exhausted' THEN 1
|
||||||
|
WHEN hard_state = 'unknown' THEN 2
|
||||||
|
WHEN probe_status = 'stale' THEN 3
|
||||||
|
ELSE 4
|
||||||
|
END ASC,
|
||||||
|
probe_failure_count DESC,
|
||||||
|
COALESCE(last_probe_success_at, 0) ASC,
|
||||||
|
COALESCE(last_scheduled_at, 0) DESC,
|
||||||
|
member_id ASC
|
||||||
|
"#,
|
||||||
|
)
|
||||||
|
.push(" LIMIT ")
|
||||||
|
.push_bind(i64_from_usize(
|
||||||
|
query.limit.max(1),
|
||||||
|
"pool probe candidate limit",
|
||||||
|
)?);
|
||||||
|
let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?;
|
||||||
|
rows.iter().map(map_score_row).collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn get_pool_member_scores_by_ids(
|
||||||
|
&self,
|
||||||
|
query: &GetPoolMemberScoresByIdsQuery,
|
||||||
|
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
|
||||||
|
if query.ids.is_empty() {
|
||||||
|
return Ok(Vec::new());
|
||||||
|
}
|
||||||
|
let mut builder = QueryBuilder::<MySql>::new(SCORE_COLUMNS);
|
||||||
|
builder.push(" WHERE id IN (");
|
||||||
|
let mut separated = builder.separated(", ");
|
||||||
|
for id in &query.ids {
|
||||||
|
separated.push_bind(id.clone());
|
||||||
|
}
|
||||||
|
separated.push_unseparated(")");
|
||||||
|
let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?;
|
||||||
|
rows.iter().map(map_score_row).collect()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl PoolMemberScoreWriteRepository for MysqlPoolMemberScoreRepository {
|
||||||
|
async fn upsert_pool_member_score(
|
||||||
|
&self,
|
||||||
|
score: UpsertPoolMemberScore,
|
||||||
|
) -> Result<StoredPoolMemberScore, DataLayerError> {
|
||||||
|
score.validate()?;
|
||||||
|
let stored = score.into_stored();
|
||||||
|
let score_reason = serde_json::to_string(&stored.score_reason)
|
||||||
|
.map_err(|err| DataLayerError::InvalidInput(err.to_string()))?;
|
||||||
|
sqlx::query(
|
||||||
|
r#"
|
||||||
|
INSERT INTO pool_member_scores (
|
||||||
|
id, pool_kind, pool_id, member_kind, member_id, capability, scope_kind, scope_id,
|
||||||
|
score, hard_state, score_version, score_reason, last_ranked_at, last_scheduled_at,
|
||||||
|
last_success_at, last_failure_at, failure_count, last_probe_attempt_at,
|
||||||
|
last_probe_success_at, last_probe_failure_at, probe_failure_count, probe_status, updated_at
|
||||||
|
) VALUES (
|
||||||
|
?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?
|
||||||
|
)
|
||||||
|
ON DUPLICATE KEY UPDATE
|
||||||
|
pool_kind = VALUES(pool_kind),
|
||||||
|
pool_id = VALUES(pool_id),
|
||||||
|
member_kind = VALUES(member_kind),
|
||||||
|
member_id = VALUES(member_id),
|
||||||
|
capability = VALUES(capability),
|
||||||
|
scope_kind = VALUES(scope_kind),
|
||||||
|
scope_id = VALUES(scope_id),
|
||||||
|
score = VALUES(score),
|
||||||
|
hard_state = VALUES(hard_state),
|
||||||
|
score_version = VALUES(score_version),
|
||||||
|
score_reason = VALUES(score_reason),
|
||||||
|
last_ranked_at = VALUES(last_ranked_at),
|
||||||
|
last_scheduled_at = COALESCE(VALUES(last_scheduled_at), last_scheduled_at),
|
||||||
|
last_success_at = COALESCE(VALUES(last_success_at), last_success_at),
|
||||||
|
last_failure_at = COALESCE(VALUES(last_failure_at), last_failure_at),
|
||||||
|
failure_count = VALUES(failure_count),
|
||||||
|
last_probe_attempt_at = COALESCE(VALUES(last_probe_attempt_at), last_probe_attempt_at),
|
||||||
|
last_probe_success_at = COALESCE(VALUES(last_probe_success_at), last_probe_success_at),
|
||||||
|
last_probe_failure_at = COALESCE(VALUES(last_probe_failure_at), last_probe_failure_at),
|
||||||
|
probe_failure_count = VALUES(probe_failure_count),
|
||||||
|
probe_status = VALUES(probe_status),
|
||||||
|
updated_at = VALUES(updated_at)
|
||||||
|
"#,
|
||||||
|
)
|
||||||
|
.bind(stored.id.as_str())
|
||||||
|
.bind(stored.pool_kind.as_str())
|
||||||
|
.bind(stored.pool_id.as_str())
|
||||||
|
.bind(stored.member_kind.as_str())
|
||||||
|
.bind(stored.member_id.as_str())
|
||||||
|
.bind(stored.capability.as_str())
|
||||||
|
.bind(stored.scope_kind.as_str())
|
||||||
|
.bind(stored.scope_id.as_deref())
|
||||||
|
.bind(stored.score)
|
||||||
|
.bind(stored.hard_state.as_database())
|
||||||
|
.bind(i64_from_u64(stored.score_version, "pool score version")?)
|
||||||
|
.bind(score_reason)
|
||||||
|
.bind(i64_opt_from_u64(
|
||||||
|
stored.last_ranked_at,
|
||||||
|
"pool score last_ranked_at",
|
||||||
|
)?)
|
||||||
|
.bind(i64_opt_from_u64(
|
||||||
|
stored.last_scheduled_at,
|
||||||
|
"pool score last_scheduled_at",
|
||||||
|
)?)
|
||||||
|
.bind(i64_opt_from_u64(
|
||||||
|
stored.last_success_at,
|
||||||
|
"pool score last_success_at",
|
||||||
|
)?)
|
||||||
|
.bind(i64_opt_from_u64(
|
||||||
|
stored.last_failure_at,
|
||||||
|
"pool score last_failure_at",
|
||||||
|
)?)
|
||||||
|
.bind(i64_from_u64(
|
||||||
|
stored.failure_count,
|
||||||
|
"pool score failure_count",
|
||||||
|
)?)
|
||||||
|
.bind(i64_opt_from_u64(
|
||||||
|
stored.last_probe_attempt_at,
|
||||||
|
"pool score last_probe_attempt_at",
|
||||||
|
)?)
|
||||||
|
.bind(i64_opt_from_u64(
|
||||||
|
stored.last_probe_success_at,
|
||||||
|
"pool score last_probe_success_at",
|
||||||
|
)?)
|
||||||
|
.bind(i64_opt_from_u64(
|
||||||
|
stored.last_probe_failure_at,
|
||||||
|
"pool score last_probe_failure_at",
|
||||||
|
)?)
|
||||||
|
.bind(i64_from_u64(
|
||||||
|
stored.probe_failure_count,
|
||||||
|
"pool score probe_failure_count",
|
||||||
|
)?)
|
||||||
|
.bind(stored.probe_status.as_database())
|
||||||
|
.bind(i64_from_u64(stored.updated_at, "pool score updated_at")?)
|
||||||
|
.execute(&self.pool)
|
||||||
|
.await
|
||||||
|
.map_sql_err()?;
|
||||||
|
Ok(stored)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn mark_pool_member_probe_in_progress(
|
||||||
|
&self,
|
||||||
|
attempt: PoolMemberProbeAttempt,
|
||||||
|
) -> Result<usize, DataLayerError> {
|
||||||
|
let rows = self
|
||||||
|
.find_scores_by_identity(&attempt.identity, attempt.scope.as_ref())
|
||||||
|
.await?;
|
||||||
|
let count = rows.len();
|
||||||
|
for mut row in rows {
|
||||||
|
row.last_probe_attempt_at = Some(attempt.attempted_at);
|
||||||
|
row.probe_status = PoolMemberProbeStatus::InProgress;
|
||||||
|
row.score_reason =
|
||||||
|
merge_score_reason_patch(row.score_reason, attempt.score_reason_patch.clone());
|
||||||
|
row.updated_at = attempt.attempted_at;
|
||||||
|
self.upsert_pool_member_score(upsert_from_stored(row))
|
||||||
|
.await?;
|
||||||
|
}
|
||||||
|
Ok(count)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn record_pool_member_probe_result(
|
||||||
|
&self,
|
||||||
|
result: PoolMemberProbeResult,
|
||||||
|
) -> Result<usize, DataLayerError> {
|
||||||
|
let rows = self
|
||||||
|
.find_scores_by_identity(&result.identity, result.scope.as_ref())
|
||||||
|
.await?;
|
||||||
|
let count = rows.len();
|
||||||
|
for mut row in rows {
|
||||||
|
row.last_probe_attempt_at = Some(result.attempted_at);
|
||||||
|
row.probe_status = result.probe_status;
|
||||||
|
if result.succeeded {
|
||||||
|
row.last_probe_success_at = Some(result.attempted_at);
|
||||||
|
row.probe_failure_count = 0;
|
||||||
|
} else {
|
||||||
|
row.last_probe_failure_at = Some(result.attempted_at);
|
||||||
|
row.probe_failure_count = row.probe_failure_count.saturating_add(1);
|
||||||
|
}
|
||||||
|
if let Some(hard_state) = result.hard_state {
|
||||||
|
row.hard_state = hard_state;
|
||||||
|
}
|
||||||
|
row.score_reason =
|
||||||
|
merge_score_reason_patch(row.score_reason, result.score_reason_patch.clone());
|
||||||
|
row.updated_at = result.attempted_at;
|
||||||
|
self.upsert_pool_member_score(upsert_from_stored(row))
|
||||||
|
.await?;
|
||||||
|
}
|
||||||
|
Ok(count)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn record_pool_member_schedule_feedback(
|
||||||
|
&self,
|
||||||
|
feedback: PoolMemberScheduleFeedback,
|
||||||
|
) -> Result<usize, DataLayerError> {
|
||||||
|
let rows = self
|
||||||
|
.find_scores_by_identity(&feedback.identity, feedback.scope.as_ref())
|
||||||
|
.await?;
|
||||||
|
let count = rows.len();
|
||||||
|
for mut row in rows {
|
||||||
|
row.last_scheduled_at = Some(feedback.scheduled_at);
|
||||||
|
match feedback.succeeded {
|
||||||
|
Some(true) => row.last_success_at = Some(feedback.scheduled_at),
|
||||||
|
Some(false) => {
|
||||||
|
row.last_failure_at = Some(feedback.scheduled_at);
|
||||||
|
row.failure_count = row.failure_count.saturating_add(1);
|
||||||
|
}
|
||||||
|
None => {}
|
||||||
|
}
|
||||||
|
if let Some(hard_state) = feedback.hard_state {
|
||||||
|
row.hard_state = hard_state;
|
||||||
|
}
|
||||||
|
row.score = score_with_delta(row.score, feedback.score_delta);
|
||||||
|
row.score_reason =
|
||||||
|
merge_score_reason_patch(row.score_reason, feedback.score_reason_patch.clone());
|
||||||
|
row.updated_at = feedback.scheduled_at;
|
||||||
|
self.upsert_pool_member_score(upsert_from_stored(row))
|
||||||
|
.await?;
|
||||||
|
}
|
||||||
|
Ok(count)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn mark_pool_member_hard_state(
|
||||||
|
&self,
|
||||||
|
identity: &PoolMemberIdentity,
|
||||||
|
scope: Option<&PoolScoreScope>,
|
||||||
|
hard_state: PoolMemberHardState,
|
||||||
|
updated_at: u64,
|
||||||
|
) -> Result<usize, DataLayerError> {
|
||||||
|
let rows = self.find_scores_by_identity(identity, scope).await?;
|
||||||
|
let count = rows.len();
|
||||||
|
for mut row in rows {
|
||||||
|
row.hard_state = hard_state;
|
||||||
|
row.updated_at = updated_at;
|
||||||
|
self.upsert_pool_member_score(upsert_from_stored(row))
|
||||||
|
.await?;
|
||||||
|
}
|
||||||
|
Ok(count)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn delete_pool_member_scores_for_member(
|
||||||
|
&self,
|
||||||
|
identity: &PoolMemberIdentity,
|
||||||
|
) -> Result<usize, DataLayerError> {
|
||||||
|
let result = sqlx::query(
|
||||||
|
r#"
|
||||||
|
DELETE FROM pool_member_scores
|
||||||
|
WHERE pool_kind = ? AND pool_id = ? AND member_kind = ? AND member_id = ?
|
||||||
|
"#,
|
||||||
|
)
|
||||||
|
.bind(identity.pool_kind.as_str())
|
||||||
|
.bind(identity.pool_id.as_str())
|
||||||
|
.bind(identity.member_kind.as_str())
|
||||||
|
.bind(identity.member_id.as_str())
|
||||||
|
.execute(&self.pool)
|
||||||
|
.await
|
||||||
|
.map_sql_err()?;
|
||||||
|
Ok(result.rows_affected() as usize)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn map_score_row(row: &MySqlRow) -> Result<StoredPoolMemberScore, DataLayerError> {
|
||||||
|
let score_reason_raw: String = row.try_get("score_reason").map_sql_err()?;
|
||||||
|
Ok(StoredPoolMemberScore {
|
||||||
|
id: row.try_get("id").map_sql_err()?,
|
||||||
|
pool_kind: row.try_get("pool_kind").map_sql_err()?,
|
||||||
|
pool_id: row.try_get("pool_id").map_sql_err()?,
|
||||||
|
member_kind: row.try_get("member_kind").map_sql_err()?,
|
||||||
|
member_id: row.try_get("member_id").map_sql_err()?,
|
||||||
|
capability: row.try_get("capability").map_sql_err()?,
|
||||||
|
scope_kind: row.try_get("scope_kind").map_sql_err()?,
|
||||||
|
scope_id: row.try_get("scope_id").map_sql_err()?,
|
||||||
|
score: row.try_get("score").map_sql_err()?,
|
||||||
|
hard_state: PoolMemberHardState::from_database(
|
||||||
|
row.try_get::<String, _>("hard_state")
|
||||||
|
.map_sql_err()?
|
||||||
|
.as_str(),
|
||||||
|
)?,
|
||||||
|
score_version: u64_from_i64(
|
||||||
|
row.try_get("score_version").map_sql_err()?,
|
||||||
|
"pool_member_scores.score_version",
|
||||||
|
)?,
|
||||||
|
score_reason: serde_json::from_str(&score_reason_raw).unwrap_or(serde_json::Value::Null),
|
||||||
|
last_ranked_at: u64_opt_from_i64(
|
||||||
|
row.try_get("last_ranked_at").map_sql_err()?,
|
||||||
|
"pool_member_scores.last_ranked_at",
|
||||||
|
)?,
|
||||||
|
last_scheduled_at: u64_opt_from_i64(
|
||||||
|
row.try_get("last_scheduled_at").map_sql_err()?,
|
||||||
|
"pool_member_scores.last_scheduled_at",
|
||||||
|
)?,
|
||||||
|
last_success_at: u64_opt_from_i64(
|
||||||
|
row.try_get("last_success_at").map_sql_err()?,
|
||||||
|
"pool_member_scores.last_success_at",
|
||||||
|
)?,
|
||||||
|
last_failure_at: u64_opt_from_i64(
|
||||||
|
row.try_get("last_failure_at").map_sql_err()?,
|
||||||
|
"pool_member_scores.last_failure_at",
|
||||||
|
)?,
|
||||||
|
failure_count: u64_from_i64(
|
||||||
|
row.try_get("failure_count").map_sql_err()?,
|
||||||
|
"pool_member_scores.failure_count",
|
||||||
|
)?,
|
||||||
|
last_probe_attempt_at: u64_opt_from_i64(
|
||||||
|
row.try_get("last_probe_attempt_at").map_sql_err()?,
|
||||||
|
"pool_member_scores.last_probe_attempt_at",
|
||||||
|
)?,
|
||||||
|
last_probe_success_at: u64_opt_from_i64(
|
||||||
|
row.try_get("last_probe_success_at").map_sql_err()?,
|
||||||
|
"pool_member_scores.last_probe_success_at",
|
||||||
|
)?,
|
||||||
|
last_probe_failure_at: u64_opt_from_i64(
|
||||||
|
row.try_get("last_probe_failure_at").map_sql_err()?,
|
||||||
|
"pool_member_scores.last_probe_failure_at",
|
||||||
|
)?,
|
||||||
|
probe_failure_count: u64_from_i64(
|
||||||
|
row.try_get("probe_failure_count").map_sql_err()?,
|
||||||
|
"pool_member_scores.probe_failure_count",
|
||||||
|
)?,
|
||||||
|
probe_status: PoolMemberProbeStatus::from_database(
|
||||||
|
row.try_get::<String, _>("probe_status")
|
||||||
|
.map_sql_err()?
|
||||||
|
.as_str(),
|
||||||
|
)?,
|
||||||
|
updated_at: u64_from_i64(
|
||||||
|
row.try_get("updated_at").map_sql_err()?,
|
||||||
|
"pool_member_scores.updated_at",
|
||||||
|
)?,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn upsert_from_stored(score: StoredPoolMemberScore) -> UpsertPoolMemberScore {
|
||||||
|
UpsertPoolMemberScore {
|
||||||
|
id: score.id,
|
||||||
|
identity: PoolMemberIdentity {
|
||||||
|
pool_kind: score.pool_kind,
|
||||||
|
pool_id: score.pool_id,
|
||||||
|
member_kind: score.member_kind,
|
||||||
|
member_id: score.member_id,
|
||||||
|
},
|
||||||
|
scope: PoolScoreScope {
|
||||||
|
capability: score.capability,
|
||||||
|
scope_kind: score.scope_kind,
|
||||||
|
scope_id: score.scope_id,
|
||||||
|
},
|
||||||
|
score: score.score,
|
||||||
|
hard_state: score.hard_state,
|
||||||
|
score_version: score.score_version,
|
||||||
|
score_reason: score.score_reason,
|
||||||
|
last_ranked_at: score.last_ranked_at,
|
||||||
|
last_scheduled_at: score.last_scheduled_at,
|
||||||
|
last_success_at: score.last_success_at,
|
||||||
|
last_failure_at: score.last_failure_at,
|
||||||
|
failure_count: score.failure_count,
|
||||||
|
last_probe_attempt_at: score.last_probe_attempt_at,
|
||||||
|
last_probe_success_at: score.last_probe_success_at,
|
||||||
|
last_probe_failure_at: score.last_probe_failure_at,
|
||||||
|
probe_failure_count: score.probe_failure_count,
|
||||||
|
probe_status: score.probe_status,
|
||||||
|
updated_at: score.updated_at,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn i64_from_usize(value: usize, field: &str) -> Result<i64, DataLayerError> {
|
||||||
|
i64::try_from(value)
|
||||||
|
.map_err(|_| DataLayerError::InvalidInput(format!("{field} exceeds signed 64-bit range")))
|
||||||
|
}
|
||||||
588
crates/aether-data/src/repository/pool_scores/postgres.rs
Normal file
588
crates/aether-data/src/repository/pool_scores/postgres.rs
Normal file
@@ -0,0 +1,588 @@
|
|||||||
|
use async_trait::async_trait;
|
||||||
|
use sqlx::{postgres::PgRow, PgPool, Postgres, QueryBuilder, Row};
|
||||||
|
|
||||||
|
use super::{
|
||||||
|
i64_from_u64, i64_opt_from_u64, score_with_delta, u64_from_i64, u64_opt_from_i64,
|
||||||
|
GetPoolMemberScoresByIdsQuery, ListPoolMemberProbeCandidatesQuery, ListPoolMemberScoresQuery,
|
||||||
|
ListRankedPoolMembersQuery, PoolMemberHardState, PoolMemberIdentity, PoolMemberProbeAttempt,
|
||||||
|
PoolMemberProbeResult, PoolMemberProbeStatus, PoolMemberScheduleFeedback,
|
||||||
|
PoolMemberScoreWriteRepository, PoolScoreReadRepository, PoolScoreScope, StoredPoolMemberScore,
|
||||||
|
UpsertPoolMemberScore,
|
||||||
|
};
|
||||||
|
use crate::error::SqlxResultExt;
|
||||||
|
use crate::repository::pool_scores::merge_score_reason_patch;
|
||||||
|
use crate::DataLayerError;
|
||||||
|
|
||||||
|
const SCORE_COLUMNS: &str = r#"
|
||||||
|
SELECT
|
||||||
|
id,
|
||||||
|
pool_kind,
|
||||||
|
pool_id,
|
||||||
|
member_kind,
|
||||||
|
member_id,
|
||||||
|
capability,
|
||||||
|
scope_kind,
|
||||||
|
scope_id,
|
||||||
|
score,
|
||||||
|
hard_state,
|
||||||
|
score_version,
|
||||||
|
score_reason,
|
||||||
|
last_ranked_at,
|
||||||
|
last_scheduled_at,
|
||||||
|
last_success_at,
|
||||||
|
last_failure_at,
|
||||||
|
failure_count,
|
||||||
|
last_probe_attempt_at,
|
||||||
|
last_probe_success_at,
|
||||||
|
last_probe_failure_at,
|
||||||
|
probe_failure_count,
|
||||||
|
probe_status,
|
||||||
|
updated_at
|
||||||
|
FROM pool_member_scores
|
||||||
|
"#;
|
||||||
|
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct PostgresPoolMemberScoreRepository {
|
||||||
|
pool: PgPool,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl PostgresPoolMemberScoreRepository {
|
||||||
|
pub fn new(pool: PgPool) -> Self {
|
||||||
|
Self { pool }
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn find_scores_by_identity(
|
||||||
|
&self,
|
||||||
|
identity: &PoolMemberIdentity,
|
||||||
|
scope: Option<&PoolScoreScope>,
|
||||||
|
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
|
||||||
|
let mut builder = QueryBuilder::<Postgres>::new(SCORE_COLUMNS);
|
||||||
|
builder
|
||||||
|
.push(" WHERE pool_kind = ")
|
||||||
|
.push_bind(identity.pool_kind.clone())
|
||||||
|
.push(" AND pool_id = ")
|
||||||
|
.push_bind(identity.pool_id.clone())
|
||||||
|
.push(" AND member_kind = ")
|
||||||
|
.push_bind(identity.member_kind.clone())
|
||||||
|
.push(" AND member_id = ")
|
||||||
|
.push_bind(identity.member_id.clone());
|
||||||
|
if let Some(scope) = scope {
|
||||||
|
builder
|
||||||
|
.push(" AND capability = ")
|
||||||
|
.push_bind(scope.capability.clone())
|
||||||
|
.push(" AND scope_kind = ")
|
||||||
|
.push_bind(scope.scope_kind.clone());
|
||||||
|
if let Some(scope_id) = &scope.scope_id {
|
||||||
|
builder.push(" AND scope_id = ").push_bind(scope_id.clone());
|
||||||
|
} else {
|
||||||
|
builder.push(" AND scope_id IS NULL");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
let rows = builder
|
||||||
|
.build()
|
||||||
|
.fetch_all(&self.pool)
|
||||||
|
.await
|
||||||
|
.map_postgres_err()?;
|
||||||
|
rows.iter().map(map_score_row).collect()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl PoolScoreReadRepository for PostgresPoolMemberScoreRepository {
|
||||||
|
async fn list_ranked_pool_members(
|
||||||
|
&self,
|
||||||
|
query: &ListRankedPoolMembersQuery,
|
||||||
|
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
|
||||||
|
let mut builder = QueryBuilder::<Postgres>::new(SCORE_COLUMNS);
|
||||||
|
builder
|
||||||
|
.push(" WHERE pool_kind = ")
|
||||||
|
.push_bind(query.pool_kind.clone())
|
||||||
|
.push(" AND pool_id = ")
|
||||||
|
.push_bind(query.pool_id.clone())
|
||||||
|
.push(" AND capability = ")
|
||||||
|
.push_bind(query.capability.clone())
|
||||||
|
.push(" AND scope_kind = ")
|
||||||
|
.push_bind(query.scope_kind.clone());
|
||||||
|
if let Some(scope_id) = &query.scope_id {
|
||||||
|
builder.push(" AND scope_id = ").push_bind(scope_id.clone());
|
||||||
|
} else {
|
||||||
|
builder.push(" AND scope_id IS NULL");
|
||||||
|
}
|
||||||
|
if !query.hard_states.is_empty() {
|
||||||
|
builder.push(" AND hard_state IN (");
|
||||||
|
let mut separated = builder.separated(", ");
|
||||||
|
for state in &query.hard_states {
|
||||||
|
separated.push_bind(state.as_database());
|
||||||
|
}
|
||||||
|
separated.push_unseparated(")");
|
||||||
|
}
|
||||||
|
if let Some(statuses) = &query.probe_statuses {
|
||||||
|
if !statuses.is_empty() {
|
||||||
|
builder.push(" AND probe_status IN (");
|
||||||
|
let mut separated = builder.separated(", ");
|
||||||
|
for status in statuses {
|
||||||
|
separated.push_bind(status.as_database());
|
||||||
|
}
|
||||||
|
separated.push_unseparated(")");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
builder
|
||||||
|
.push(" ORDER BY score DESC, last_ranked_at DESC NULLS LAST, member_id ASC, id ASC")
|
||||||
|
.push(" LIMIT ")
|
||||||
|
.push_bind(i64_from_usize(query.limit.max(1), "pool score limit")?)
|
||||||
|
.push(" OFFSET ")
|
||||||
|
.push_bind(i64_from_usize(query.offset, "pool score offset")?);
|
||||||
|
let rows = builder
|
||||||
|
.build()
|
||||||
|
.fetch_all(&self.pool)
|
||||||
|
.await
|
||||||
|
.map_postgres_err()?;
|
||||||
|
rows.iter().map(map_score_row).collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn list_pool_member_scores(
|
||||||
|
&self,
|
||||||
|
query: &ListPoolMemberScoresQuery,
|
||||||
|
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
|
||||||
|
let mut builder = QueryBuilder::<Postgres>::new(SCORE_COLUMNS);
|
||||||
|
builder
|
||||||
|
.push(" WHERE pool_kind = ")
|
||||||
|
.push_bind(query.pool_kind.clone())
|
||||||
|
.push(" AND pool_id = ")
|
||||||
|
.push_bind(query.pool_id.clone());
|
||||||
|
if let Some(capability) = &query.capability {
|
||||||
|
builder
|
||||||
|
.push(" AND capability = ")
|
||||||
|
.push_bind(capability.clone());
|
||||||
|
}
|
||||||
|
if let Some(scope_kind) = &query.scope_kind {
|
||||||
|
builder
|
||||||
|
.push(" AND scope_kind = ")
|
||||||
|
.push_bind(scope_kind.clone());
|
||||||
|
}
|
||||||
|
if let Some(scope_id) = &query.scope_id {
|
||||||
|
builder.push(" AND scope_id = ").push_bind(scope_id.clone());
|
||||||
|
}
|
||||||
|
if !query.hard_states.is_empty() {
|
||||||
|
builder.push(" AND hard_state IN (");
|
||||||
|
let mut separated = builder.separated(", ");
|
||||||
|
for state in &query.hard_states {
|
||||||
|
separated.push_bind(state.as_database());
|
||||||
|
}
|
||||||
|
separated.push_unseparated(")");
|
||||||
|
}
|
||||||
|
if let Some(statuses) = &query.probe_statuses {
|
||||||
|
if !statuses.is_empty() {
|
||||||
|
builder.push(" AND probe_status IN (");
|
||||||
|
let mut separated = builder.separated(", ");
|
||||||
|
for status in statuses {
|
||||||
|
separated.push_bind(status.as_database());
|
||||||
|
}
|
||||||
|
separated.push_unseparated(")");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
builder
|
||||||
|
.push(" ORDER BY score DESC, last_ranked_at DESC NULLS LAST, member_id ASC, id ASC")
|
||||||
|
.push(" LIMIT ")
|
||||||
|
.push_bind(i64_from_usize(query.limit.max(1), "pool score limit")?)
|
||||||
|
.push(" OFFSET ")
|
||||||
|
.push_bind(i64_from_usize(query.offset, "pool score offset")?);
|
||||||
|
let rows = builder
|
||||||
|
.build()
|
||||||
|
.fetch_all(&self.pool)
|
||||||
|
.await
|
||||||
|
.map_postgres_err()?;
|
||||||
|
rows.iter().map(map_score_row).collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn list_pool_member_probe_candidates(
|
||||||
|
&self,
|
||||||
|
query: &ListPoolMemberProbeCandidatesQuery,
|
||||||
|
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
|
||||||
|
let mut builder = QueryBuilder::<Postgres>::new(SCORE_COLUMNS);
|
||||||
|
builder
|
||||||
|
.push(" WHERE pool_kind = ")
|
||||||
|
.push_bind(query.pool_kind.clone())
|
||||||
|
.push(" AND pool_id = ")
|
||||||
|
.push_bind(query.pool_id.clone());
|
||||||
|
if let Some(capability) = &query.capability {
|
||||||
|
builder
|
||||||
|
.push(" AND capability = ")
|
||||||
|
.push_bind(capability.clone());
|
||||||
|
}
|
||||||
|
builder
|
||||||
|
.push(" AND hard_state IN ('available','unknown','cooldown','quota_exhausted')")
|
||||||
|
.push(" AND (probe_status IN ('never','failed','stale')")
|
||||||
|
.push(" OR (probe_status = 'ok' AND (last_probe_success_at IS NULL OR last_probe_success_at <= ")
|
||||||
|
.push_bind(i64_from_u64(
|
||||||
|
query.stale_before_unix_secs,
|
||||||
|
"pool probe stale_before_unix_secs",
|
||||||
|
)?)
|
||||||
|
.push("))")
|
||||||
|
.push(" OR (probe_status = 'in_progress' AND (last_probe_attempt_at IS NULL OR last_probe_attempt_at <= ")
|
||||||
|
.push_bind(i64_from_u64(
|
||||||
|
query.stale_before_unix_secs,
|
||||||
|
"pool probe stale_before_unix_secs",
|
||||||
|
)?)
|
||||||
|
.push(")))")
|
||||||
|
.push(
|
||||||
|
r#"
|
||||||
|
ORDER BY
|
||||||
|
CASE
|
||||||
|
WHEN last_scheduled_at IS NOT NULL AND probe_status <> 'ok' THEN 0
|
||||||
|
WHEN hard_state = 'quota_exhausted' THEN 1
|
||||||
|
WHEN hard_state = 'unknown' THEN 2
|
||||||
|
WHEN probe_status = 'stale' THEN 3
|
||||||
|
ELSE 4
|
||||||
|
END ASC,
|
||||||
|
probe_failure_count DESC,
|
||||||
|
COALESCE(last_probe_success_at, 0) ASC,
|
||||||
|
COALESCE(last_scheduled_at, 0) DESC,
|
||||||
|
member_id ASC
|
||||||
|
"#,
|
||||||
|
)
|
||||||
|
.push(" LIMIT ")
|
||||||
|
.push_bind(i64_from_usize(
|
||||||
|
query.limit.max(1),
|
||||||
|
"pool probe candidate limit",
|
||||||
|
)?);
|
||||||
|
let rows = builder
|
||||||
|
.build()
|
||||||
|
.fetch_all(&self.pool)
|
||||||
|
.await
|
||||||
|
.map_postgres_err()?;
|
||||||
|
rows.iter().map(map_score_row).collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn get_pool_member_scores_by_ids(
|
||||||
|
&self,
|
||||||
|
query: &GetPoolMemberScoresByIdsQuery,
|
||||||
|
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
|
||||||
|
if query.ids.is_empty() {
|
||||||
|
return Ok(Vec::new());
|
||||||
|
}
|
||||||
|
let mut builder = QueryBuilder::<Postgres>::new(SCORE_COLUMNS);
|
||||||
|
builder.push(" WHERE id IN (");
|
||||||
|
let mut separated = builder.separated(", ");
|
||||||
|
for id in &query.ids {
|
||||||
|
separated.push_bind(id.clone());
|
||||||
|
}
|
||||||
|
separated.push_unseparated(")");
|
||||||
|
let rows = builder
|
||||||
|
.build()
|
||||||
|
.fetch_all(&self.pool)
|
||||||
|
.await
|
||||||
|
.map_postgres_err()?;
|
||||||
|
rows.iter().map(map_score_row).collect()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl PoolMemberScoreWriteRepository for PostgresPoolMemberScoreRepository {
|
||||||
|
async fn upsert_pool_member_score(
|
||||||
|
&self,
|
||||||
|
score: UpsertPoolMemberScore,
|
||||||
|
) -> Result<StoredPoolMemberScore, DataLayerError> {
|
||||||
|
score.validate()?;
|
||||||
|
let stored = score.clone().into_stored();
|
||||||
|
sqlx::query(
|
||||||
|
r#"
|
||||||
|
INSERT INTO pool_member_scores (
|
||||||
|
id, pool_kind, pool_id, member_kind, member_id, capability, scope_kind, scope_id,
|
||||||
|
score, hard_state, score_version, score_reason, last_ranked_at, last_scheduled_at,
|
||||||
|
last_success_at, last_failure_at, failure_count, last_probe_attempt_at,
|
||||||
|
last_probe_success_at, last_probe_failure_at, probe_failure_count, probe_status, updated_at
|
||||||
|
) VALUES (
|
||||||
|
$1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17,$18,$19,$20,$21,$22,$23
|
||||||
|
)
|
||||||
|
ON CONFLICT(id) DO UPDATE SET
|
||||||
|
pool_kind = EXCLUDED.pool_kind,
|
||||||
|
pool_id = EXCLUDED.pool_id,
|
||||||
|
member_kind = EXCLUDED.member_kind,
|
||||||
|
member_id = EXCLUDED.member_id,
|
||||||
|
capability = EXCLUDED.capability,
|
||||||
|
scope_kind = EXCLUDED.scope_kind,
|
||||||
|
scope_id = EXCLUDED.scope_id,
|
||||||
|
score = EXCLUDED.score,
|
||||||
|
hard_state = EXCLUDED.hard_state,
|
||||||
|
score_version = EXCLUDED.score_version,
|
||||||
|
score_reason = EXCLUDED.score_reason,
|
||||||
|
last_ranked_at = EXCLUDED.last_ranked_at,
|
||||||
|
last_scheduled_at = COALESCE(EXCLUDED.last_scheduled_at, pool_member_scores.last_scheduled_at),
|
||||||
|
last_success_at = COALESCE(EXCLUDED.last_success_at, pool_member_scores.last_success_at),
|
||||||
|
last_failure_at = COALESCE(EXCLUDED.last_failure_at, pool_member_scores.last_failure_at),
|
||||||
|
failure_count = EXCLUDED.failure_count,
|
||||||
|
last_probe_attempt_at = COALESCE(EXCLUDED.last_probe_attempt_at, pool_member_scores.last_probe_attempt_at),
|
||||||
|
last_probe_success_at = COALESCE(EXCLUDED.last_probe_success_at, pool_member_scores.last_probe_success_at),
|
||||||
|
last_probe_failure_at = COALESCE(EXCLUDED.last_probe_failure_at, pool_member_scores.last_probe_failure_at),
|
||||||
|
probe_failure_count = EXCLUDED.probe_failure_count,
|
||||||
|
probe_status = EXCLUDED.probe_status,
|
||||||
|
updated_at = EXCLUDED.updated_at
|
||||||
|
"#,
|
||||||
|
)
|
||||||
|
.bind(stored.id.as_str())
|
||||||
|
.bind(stored.pool_kind.as_str())
|
||||||
|
.bind(stored.pool_id.as_str())
|
||||||
|
.bind(stored.member_kind.as_str())
|
||||||
|
.bind(stored.member_id.as_str())
|
||||||
|
.bind(stored.capability.as_str())
|
||||||
|
.bind(stored.scope_kind.as_str())
|
||||||
|
.bind(stored.scope_id.as_deref())
|
||||||
|
.bind(stored.score)
|
||||||
|
.bind(stored.hard_state.as_database())
|
||||||
|
.bind(i64_from_u64(stored.score_version, "pool score version")?)
|
||||||
|
.bind(&stored.score_reason)
|
||||||
|
.bind(i64_opt_from_u64(stored.last_ranked_at, "pool score last_ranked_at")?)
|
||||||
|
.bind(i64_opt_from_u64(stored.last_scheduled_at, "pool score last_scheduled_at")?)
|
||||||
|
.bind(i64_opt_from_u64(stored.last_success_at, "pool score last_success_at")?)
|
||||||
|
.bind(i64_opt_from_u64(stored.last_failure_at, "pool score last_failure_at")?)
|
||||||
|
.bind(i64_from_u64(stored.failure_count, "pool score failure_count")?)
|
||||||
|
.bind(i64_opt_from_u64(
|
||||||
|
stored.last_probe_attempt_at,
|
||||||
|
"pool score last_probe_attempt_at",
|
||||||
|
)?)
|
||||||
|
.bind(i64_opt_from_u64(
|
||||||
|
stored.last_probe_success_at,
|
||||||
|
"pool score last_probe_success_at",
|
||||||
|
)?)
|
||||||
|
.bind(i64_opt_from_u64(
|
||||||
|
stored.last_probe_failure_at,
|
||||||
|
"pool score last_probe_failure_at",
|
||||||
|
)?)
|
||||||
|
.bind(i64_from_u64(
|
||||||
|
stored.probe_failure_count,
|
||||||
|
"pool score probe_failure_count",
|
||||||
|
)?)
|
||||||
|
.bind(stored.probe_status.as_database())
|
||||||
|
.bind(i64_from_u64(stored.updated_at, "pool score updated_at")?)
|
||||||
|
.execute(&self.pool)
|
||||||
|
.await
|
||||||
|
.map_postgres_err()?;
|
||||||
|
Ok(stored)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn mark_pool_member_probe_in_progress(
|
||||||
|
&self,
|
||||||
|
attempt: PoolMemberProbeAttempt,
|
||||||
|
) -> Result<usize, DataLayerError> {
|
||||||
|
let rows = self
|
||||||
|
.find_scores_by_identity(&attempt.identity, attempt.scope.as_ref())
|
||||||
|
.await?;
|
||||||
|
let count = rows.len();
|
||||||
|
for mut row in rows {
|
||||||
|
row.last_probe_attempt_at = Some(attempt.attempted_at);
|
||||||
|
row.probe_status = PoolMemberProbeStatus::InProgress;
|
||||||
|
row.score_reason =
|
||||||
|
merge_score_reason_patch(row.score_reason, attempt.score_reason_patch.clone());
|
||||||
|
row.updated_at = attempt.attempted_at;
|
||||||
|
self.upsert_pool_member_score(upsert_from_stored(row))
|
||||||
|
.await?;
|
||||||
|
}
|
||||||
|
Ok(count)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn record_pool_member_probe_result(
|
||||||
|
&self,
|
||||||
|
result: PoolMemberProbeResult,
|
||||||
|
) -> Result<usize, DataLayerError> {
|
||||||
|
let rows = self
|
||||||
|
.find_scores_by_identity(&result.identity, result.scope.as_ref())
|
||||||
|
.await?;
|
||||||
|
let count = rows.len();
|
||||||
|
for mut row in rows {
|
||||||
|
row.last_probe_attempt_at = Some(result.attempted_at);
|
||||||
|
row.probe_status = result.probe_status;
|
||||||
|
if result.succeeded {
|
||||||
|
row.last_probe_success_at = Some(result.attempted_at);
|
||||||
|
row.probe_failure_count = 0;
|
||||||
|
} else {
|
||||||
|
row.last_probe_failure_at = Some(result.attempted_at);
|
||||||
|
row.probe_failure_count = row.probe_failure_count.saturating_add(1);
|
||||||
|
}
|
||||||
|
if let Some(hard_state) = result.hard_state {
|
||||||
|
row.hard_state = hard_state;
|
||||||
|
}
|
||||||
|
row.score_reason =
|
||||||
|
merge_score_reason_patch(row.score_reason, result.score_reason_patch.clone());
|
||||||
|
row.updated_at = result.attempted_at;
|
||||||
|
self.upsert_pool_member_score(upsert_from_stored(row))
|
||||||
|
.await?;
|
||||||
|
}
|
||||||
|
Ok(count)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn record_pool_member_schedule_feedback(
|
||||||
|
&self,
|
||||||
|
feedback: PoolMemberScheduleFeedback,
|
||||||
|
) -> Result<usize, DataLayerError> {
|
||||||
|
let rows = self
|
||||||
|
.find_scores_by_identity(&feedback.identity, feedback.scope.as_ref())
|
||||||
|
.await?;
|
||||||
|
let count = rows.len();
|
||||||
|
for mut row in rows {
|
||||||
|
row.last_scheduled_at = Some(feedback.scheduled_at);
|
||||||
|
match feedback.succeeded {
|
||||||
|
Some(true) => row.last_success_at = Some(feedback.scheduled_at),
|
||||||
|
Some(false) => {
|
||||||
|
row.last_failure_at = Some(feedback.scheduled_at);
|
||||||
|
row.failure_count = row.failure_count.saturating_add(1);
|
||||||
|
}
|
||||||
|
None => {}
|
||||||
|
}
|
||||||
|
if let Some(hard_state) = feedback.hard_state {
|
||||||
|
row.hard_state = hard_state;
|
||||||
|
}
|
||||||
|
row.score = score_with_delta(row.score, feedback.score_delta);
|
||||||
|
row.score_reason =
|
||||||
|
merge_score_reason_patch(row.score_reason, feedback.score_reason_patch.clone());
|
||||||
|
row.updated_at = feedback.scheduled_at;
|
||||||
|
self.upsert_pool_member_score(upsert_from_stored(row))
|
||||||
|
.await?;
|
||||||
|
}
|
||||||
|
Ok(count)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn mark_pool_member_hard_state(
|
||||||
|
&self,
|
||||||
|
identity: &PoolMemberIdentity,
|
||||||
|
scope: Option<&PoolScoreScope>,
|
||||||
|
hard_state: PoolMemberHardState,
|
||||||
|
updated_at: u64,
|
||||||
|
) -> Result<usize, DataLayerError> {
|
||||||
|
let rows = self.find_scores_by_identity(identity, scope).await?;
|
||||||
|
let count = rows.len();
|
||||||
|
for mut row in rows {
|
||||||
|
row.hard_state = hard_state;
|
||||||
|
row.updated_at = updated_at;
|
||||||
|
self.upsert_pool_member_score(upsert_from_stored(row))
|
||||||
|
.await?;
|
||||||
|
}
|
||||||
|
Ok(count)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn delete_pool_member_scores_for_member(
|
||||||
|
&self,
|
||||||
|
identity: &PoolMemberIdentity,
|
||||||
|
) -> Result<usize, DataLayerError> {
|
||||||
|
let result = sqlx::query(
|
||||||
|
r#"
|
||||||
|
DELETE FROM pool_member_scores
|
||||||
|
WHERE pool_kind = $1 AND pool_id = $2 AND member_kind = $3 AND member_id = $4
|
||||||
|
"#,
|
||||||
|
)
|
||||||
|
.bind(identity.pool_kind.as_str())
|
||||||
|
.bind(identity.pool_id.as_str())
|
||||||
|
.bind(identity.member_kind.as_str())
|
||||||
|
.bind(identity.member_id.as_str())
|
||||||
|
.execute(&self.pool)
|
||||||
|
.await
|
||||||
|
.map_postgres_err()?;
|
||||||
|
Ok(result.rows_affected() as usize)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn map_score_row(row: &PgRow) -> Result<StoredPoolMemberScore, DataLayerError> {
|
||||||
|
Ok(StoredPoolMemberScore {
|
||||||
|
id: row.try_get("id").map_postgres_err()?,
|
||||||
|
pool_kind: row.try_get("pool_kind").map_postgres_err()?,
|
||||||
|
pool_id: row.try_get("pool_id").map_postgres_err()?,
|
||||||
|
member_kind: row.try_get("member_kind").map_postgres_err()?,
|
||||||
|
member_id: row.try_get("member_id").map_postgres_err()?,
|
||||||
|
capability: row.try_get("capability").map_postgres_err()?,
|
||||||
|
scope_kind: row.try_get("scope_kind").map_postgres_err()?,
|
||||||
|
scope_id: row.try_get("scope_id").map_postgres_err()?,
|
||||||
|
score: row.try_get("score").map_postgres_err()?,
|
||||||
|
hard_state: PoolMemberHardState::from_database(
|
||||||
|
row.try_get::<String, _>("hard_state")
|
||||||
|
.map_postgres_err()?
|
||||||
|
.as_str(),
|
||||||
|
)?,
|
||||||
|
score_version: u64_from_i64(
|
||||||
|
row.try_get("score_version").map_postgres_err()?,
|
||||||
|
"pool_member_scores.score_version",
|
||||||
|
)?,
|
||||||
|
score_reason: row.try_get("score_reason").map_postgres_err()?,
|
||||||
|
last_ranked_at: u64_opt_from_i64(
|
||||||
|
row.try_get("last_ranked_at").map_postgres_err()?,
|
||||||
|
"pool_member_scores.last_ranked_at",
|
||||||
|
)?,
|
||||||
|
last_scheduled_at: u64_opt_from_i64(
|
||||||
|
row.try_get("last_scheduled_at").map_postgres_err()?,
|
||||||
|
"pool_member_scores.last_scheduled_at",
|
||||||
|
)?,
|
||||||
|
last_success_at: u64_opt_from_i64(
|
||||||
|
row.try_get("last_success_at").map_postgres_err()?,
|
||||||
|
"pool_member_scores.last_success_at",
|
||||||
|
)?,
|
||||||
|
last_failure_at: u64_opt_from_i64(
|
||||||
|
row.try_get("last_failure_at").map_postgres_err()?,
|
||||||
|
"pool_member_scores.last_failure_at",
|
||||||
|
)?,
|
||||||
|
failure_count: u64_from_i64(
|
||||||
|
row.try_get("failure_count").map_postgres_err()?,
|
||||||
|
"pool_member_scores.failure_count",
|
||||||
|
)?,
|
||||||
|
last_probe_attempt_at: u64_opt_from_i64(
|
||||||
|
row.try_get("last_probe_attempt_at").map_postgres_err()?,
|
||||||
|
"pool_member_scores.last_probe_attempt_at",
|
||||||
|
)?,
|
||||||
|
last_probe_success_at: u64_opt_from_i64(
|
||||||
|
row.try_get("last_probe_success_at").map_postgres_err()?,
|
||||||
|
"pool_member_scores.last_probe_success_at",
|
||||||
|
)?,
|
||||||
|
last_probe_failure_at: u64_opt_from_i64(
|
||||||
|
row.try_get("last_probe_failure_at").map_postgres_err()?,
|
||||||
|
"pool_member_scores.last_probe_failure_at",
|
||||||
|
)?,
|
||||||
|
probe_failure_count: u64_from_i64(
|
||||||
|
row.try_get("probe_failure_count").map_postgres_err()?,
|
||||||
|
"pool_member_scores.probe_failure_count",
|
||||||
|
)?,
|
||||||
|
probe_status: PoolMemberProbeStatus::from_database(
|
||||||
|
row.try_get::<String, _>("probe_status")
|
||||||
|
.map_postgres_err()?
|
||||||
|
.as_str(),
|
||||||
|
)?,
|
||||||
|
updated_at: u64_from_i64(
|
||||||
|
row.try_get("updated_at").map_postgres_err()?,
|
||||||
|
"pool_member_scores.updated_at",
|
||||||
|
)?,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn upsert_from_stored(score: StoredPoolMemberScore) -> UpsertPoolMemberScore {
|
||||||
|
UpsertPoolMemberScore {
|
||||||
|
id: score.id,
|
||||||
|
identity: PoolMemberIdentity {
|
||||||
|
pool_kind: score.pool_kind,
|
||||||
|
pool_id: score.pool_id,
|
||||||
|
member_kind: score.member_kind,
|
||||||
|
member_id: score.member_id,
|
||||||
|
},
|
||||||
|
scope: PoolScoreScope {
|
||||||
|
capability: score.capability,
|
||||||
|
scope_kind: score.scope_kind,
|
||||||
|
scope_id: score.scope_id,
|
||||||
|
},
|
||||||
|
score: score.score,
|
||||||
|
hard_state: score.hard_state,
|
||||||
|
score_version: score.score_version,
|
||||||
|
score_reason: score.score_reason,
|
||||||
|
last_ranked_at: score.last_ranked_at,
|
||||||
|
last_scheduled_at: score.last_scheduled_at,
|
||||||
|
last_success_at: score.last_success_at,
|
||||||
|
last_failure_at: score.last_failure_at,
|
||||||
|
failure_count: score.failure_count,
|
||||||
|
last_probe_attempt_at: score.last_probe_attempt_at,
|
||||||
|
last_probe_success_at: score.last_probe_success_at,
|
||||||
|
last_probe_failure_at: score.last_probe_failure_at,
|
||||||
|
probe_failure_count: score.probe_failure_count,
|
||||||
|
probe_status: score.probe_status,
|
||||||
|
updated_at: score.updated_at,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn i64_from_usize(value: usize, field: &str) -> Result<i64, DataLayerError> {
|
||||||
|
i64::try_from(value)
|
||||||
|
.map_err(|_| DataLayerError::InvalidInput(format!("{field} exceeds signed 64-bit range")))
|
||||||
|
}
|
||||||
571
crates/aether-data/src/repository/pool_scores/sqlite.rs
Normal file
571
crates/aether-data/src/repository/pool_scores/sqlite.rs
Normal file
@@ -0,0 +1,571 @@
|
|||||||
|
use async_trait::async_trait;
|
||||||
|
use sqlx::{sqlite::SqliteRow, QueryBuilder, Row, Sqlite, SqlitePool};
|
||||||
|
|
||||||
|
use super::{
|
||||||
|
i64_from_u64, i64_opt_from_u64, score_with_delta, u64_from_i64, u64_opt_from_i64,
|
||||||
|
GetPoolMemberScoresByIdsQuery, ListPoolMemberProbeCandidatesQuery, ListPoolMemberScoresQuery,
|
||||||
|
ListRankedPoolMembersQuery, PoolMemberHardState, PoolMemberIdentity, PoolMemberProbeAttempt,
|
||||||
|
PoolMemberProbeResult, PoolMemberProbeStatus, PoolMemberScheduleFeedback,
|
||||||
|
PoolMemberScoreWriteRepository, PoolScoreReadRepository, PoolScoreScope, StoredPoolMemberScore,
|
||||||
|
UpsertPoolMemberScore,
|
||||||
|
};
|
||||||
|
use crate::error::SqlResultExt;
|
||||||
|
use crate::repository::pool_scores::merge_score_reason_patch;
|
||||||
|
use crate::DataLayerError;
|
||||||
|
|
||||||
|
const SCORE_COLUMNS: &str = r#"
|
||||||
|
SELECT
|
||||||
|
id,
|
||||||
|
pool_kind,
|
||||||
|
pool_id,
|
||||||
|
member_kind,
|
||||||
|
member_id,
|
||||||
|
capability,
|
||||||
|
scope_kind,
|
||||||
|
scope_id,
|
||||||
|
score,
|
||||||
|
hard_state,
|
||||||
|
score_version,
|
||||||
|
score_reason,
|
||||||
|
last_ranked_at,
|
||||||
|
last_scheduled_at,
|
||||||
|
last_success_at,
|
||||||
|
last_failure_at,
|
||||||
|
failure_count,
|
||||||
|
last_probe_attempt_at,
|
||||||
|
last_probe_success_at,
|
||||||
|
last_probe_failure_at,
|
||||||
|
probe_failure_count,
|
||||||
|
probe_status,
|
||||||
|
updated_at
|
||||||
|
FROM pool_member_scores
|
||||||
|
"#;
|
||||||
|
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct SqlitePoolMemberScoreRepository {
|
||||||
|
pool: SqlitePool,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl SqlitePoolMemberScoreRepository {
|
||||||
|
pub fn new(pool: SqlitePool) -> Self {
|
||||||
|
Self { pool }
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn find_scores_by_identity(
|
||||||
|
&self,
|
||||||
|
identity: &PoolMemberIdentity,
|
||||||
|
scope: Option<&PoolScoreScope>,
|
||||||
|
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
|
||||||
|
let mut builder = QueryBuilder::<Sqlite>::new(SCORE_COLUMNS);
|
||||||
|
builder
|
||||||
|
.push(" WHERE pool_kind = ")
|
||||||
|
.push_bind(identity.pool_kind.clone())
|
||||||
|
.push(" AND pool_id = ")
|
||||||
|
.push_bind(identity.pool_id.clone())
|
||||||
|
.push(" AND member_kind = ")
|
||||||
|
.push_bind(identity.member_kind.clone())
|
||||||
|
.push(" AND member_id = ")
|
||||||
|
.push_bind(identity.member_id.clone());
|
||||||
|
if let Some(scope) = scope {
|
||||||
|
builder
|
||||||
|
.push(" AND capability = ")
|
||||||
|
.push_bind(scope.capability.clone())
|
||||||
|
.push(" AND scope_kind = ")
|
||||||
|
.push_bind(scope.scope_kind.clone());
|
||||||
|
if let Some(scope_id) = &scope.scope_id {
|
||||||
|
builder.push(" AND scope_id = ").push_bind(scope_id.clone());
|
||||||
|
} else {
|
||||||
|
builder.push(" AND scope_id IS NULL");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?;
|
||||||
|
rows.iter().map(map_score_row).collect()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl PoolScoreReadRepository for SqlitePoolMemberScoreRepository {
|
||||||
|
async fn list_ranked_pool_members(
|
||||||
|
&self,
|
||||||
|
query: &ListRankedPoolMembersQuery,
|
||||||
|
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
|
||||||
|
let mut builder = QueryBuilder::<Sqlite>::new(SCORE_COLUMNS);
|
||||||
|
builder
|
||||||
|
.push(" WHERE pool_kind = ")
|
||||||
|
.push_bind(query.pool_kind.clone())
|
||||||
|
.push(" AND pool_id = ")
|
||||||
|
.push_bind(query.pool_id.clone())
|
||||||
|
.push(" AND capability = ")
|
||||||
|
.push_bind(query.capability.clone())
|
||||||
|
.push(" AND scope_kind = ")
|
||||||
|
.push_bind(query.scope_kind.clone());
|
||||||
|
if let Some(scope_id) = &query.scope_id {
|
||||||
|
builder.push(" AND scope_id = ").push_bind(scope_id.clone());
|
||||||
|
} else {
|
||||||
|
builder.push(" AND scope_id IS NULL");
|
||||||
|
}
|
||||||
|
if !query.hard_states.is_empty() {
|
||||||
|
builder.push(" AND hard_state IN (");
|
||||||
|
let mut separated = builder.separated(", ");
|
||||||
|
for state in &query.hard_states {
|
||||||
|
separated.push_bind(state.as_database());
|
||||||
|
}
|
||||||
|
separated.push_unseparated(")");
|
||||||
|
}
|
||||||
|
if let Some(statuses) = &query.probe_statuses {
|
||||||
|
if !statuses.is_empty() {
|
||||||
|
builder.push(" AND probe_status IN (");
|
||||||
|
let mut separated = builder.separated(", ");
|
||||||
|
for status in statuses {
|
||||||
|
separated.push_bind(status.as_database());
|
||||||
|
}
|
||||||
|
separated.push_unseparated(")");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
builder
|
||||||
|
.push(" ORDER BY score DESC, last_ranked_at DESC, member_id ASC, id ASC")
|
||||||
|
.push(" LIMIT ")
|
||||||
|
.push_bind(i64_from_usize(query.limit.max(1), "pool score limit")?)
|
||||||
|
.push(" OFFSET ")
|
||||||
|
.push_bind(i64_from_usize(query.offset, "pool score offset")?);
|
||||||
|
let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?;
|
||||||
|
rows.iter().map(map_score_row).collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn list_pool_member_scores(
|
||||||
|
&self,
|
||||||
|
query: &ListPoolMemberScoresQuery,
|
||||||
|
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
|
||||||
|
let mut builder = QueryBuilder::<Sqlite>::new(SCORE_COLUMNS);
|
||||||
|
builder
|
||||||
|
.push(" WHERE pool_kind = ")
|
||||||
|
.push_bind(query.pool_kind.clone())
|
||||||
|
.push(" AND pool_id = ")
|
||||||
|
.push_bind(query.pool_id.clone());
|
||||||
|
if let Some(capability) = &query.capability {
|
||||||
|
builder
|
||||||
|
.push(" AND capability = ")
|
||||||
|
.push_bind(capability.clone());
|
||||||
|
}
|
||||||
|
if let Some(scope_kind) = &query.scope_kind {
|
||||||
|
builder
|
||||||
|
.push(" AND scope_kind = ")
|
||||||
|
.push_bind(scope_kind.clone());
|
||||||
|
}
|
||||||
|
if let Some(scope_id) = &query.scope_id {
|
||||||
|
builder.push(" AND scope_id = ").push_bind(scope_id.clone());
|
||||||
|
}
|
||||||
|
if !query.hard_states.is_empty() {
|
||||||
|
builder.push(" AND hard_state IN (");
|
||||||
|
let mut separated = builder.separated(", ");
|
||||||
|
for state in &query.hard_states {
|
||||||
|
separated.push_bind(state.as_database());
|
||||||
|
}
|
||||||
|
separated.push_unseparated(")");
|
||||||
|
}
|
||||||
|
if let Some(statuses) = &query.probe_statuses {
|
||||||
|
if !statuses.is_empty() {
|
||||||
|
builder.push(" AND probe_status IN (");
|
||||||
|
let mut separated = builder.separated(", ");
|
||||||
|
for status in statuses {
|
||||||
|
separated.push_bind(status.as_database());
|
||||||
|
}
|
||||||
|
separated.push_unseparated(")");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
builder
|
||||||
|
.push(" ORDER BY score DESC, last_ranked_at DESC, member_id ASC, id ASC")
|
||||||
|
.push(" LIMIT ")
|
||||||
|
.push_bind(i64_from_usize(query.limit.max(1), "pool score limit")?)
|
||||||
|
.push(" OFFSET ")
|
||||||
|
.push_bind(i64_from_usize(query.offset, "pool score offset")?);
|
||||||
|
let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?;
|
||||||
|
rows.iter().map(map_score_row).collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn list_pool_member_probe_candidates(
|
||||||
|
&self,
|
||||||
|
query: &ListPoolMemberProbeCandidatesQuery,
|
||||||
|
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
|
||||||
|
let mut builder = QueryBuilder::<Sqlite>::new(SCORE_COLUMNS);
|
||||||
|
builder
|
||||||
|
.push(" WHERE pool_kind = ")
|
||||||
|
.push_bind(query.pool_kind.clone())
|
||||||
|
.push(" AND pool_id = ")
|
||||||
|
.push_bind(query.pool_id.clone());
|
||||||
|
if let Some(capability) = &query.capability {
|
||||||
|
builder
|
||||||
|
.push(" AND capability = ")
|
||||||
|
.push_bind(capability.clone());
|
||||||
|
}
|
||||||
|
builder
|
||||||
|
.push(" AND hard_state IN ('available','unknown','cooldown','quota_exhausted')")
|
||||||
|
.push(" AND (probe_status IN ('never','failed','stale')")
|
||||||
|
.push(" OR (probe_status = 'ok' AND (last_probe_success_at IS NULL OR last_probe_success_at <= ")
|
||||||
|
.push_bind(i64_from_u64(
|
||||||
|
query.stale_before_unix_secs,
|
||||||
|
"pool probe stale_before_unix_secs",
|
||||||
|
)?)
|
||||||
|
.push("))")
|
||||||
|
.push(" OR (probe_status = 'in_progress' AND (last_probe_attempt_at IS NULL OR last_probe_attempt_at <= ")
|
||||||
|
.push_bind(i64_from_u64(
|
||||||
|
query.stale_before_unix_secs,
|
||||||
|
"pool probe stale_before_unix_secs",
|
||||||
|
)?)
|
||||||
|
.push(")))")
|
||||||
|
.push(
|
||||||
|
r#"
|
||||||
|
ORDER BY
|
||||||
|
CASE
|
||||||
|
WHEN last_scheduled_at IS NOT NULL AND probe_status <> 'ok' THEN 0
|
||||||
|
WHEN hard_state = 'quota_exhausted' THEN 1
|
||||||
|
WHEN hard_state = 'unknown' THEN 2
|
||||||
|
WHEN probe_status = 'stale' THEN 3
|
||||||
|
ELSE 4
|
||||||
|
END ASC,
|
||||||
|
probe_failure_count DESC,
|
||||||
|
COALESCE(last_probe_success_at, 0) ASC,
|
||||||
|
COALESCE(last_scheduled_at, 0) DESC,
|
||||||
|
member_id ASC
|
||||||
|
"#,
|
||||||
|
)
|
||||||
|
.push(" LIMIT ")
|
||||||
|
.push_bind(i64_from_usize(
|
||||||
|
query.limit.max(1),
|
||||||
|
"pool probe candidate limit",
|
||||||
|
)?);
|
||||||
|
let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?;
|
||||||
|
rows.iter().map(map_score_row).collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn get_pool_member_scores_by_ids(
|
||||||
|
&self,
|
||||||
|
query: &GetPoolMemberScoresByIdsQuery,
|
||||||
|
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
|
||||||
|
if query.ids.is_empty() {
|
||||||
|
return Ok(Vec::new());
|
||||||
|
}
|
||||||
|
let mut builder = QueryBuilder::<Sqlite>::new(SCORE_COLUMNS);
|
||||||
|
builder.push(" WHERE id IN (");
|
||||||
|
let mut separated = builder.separated(", ");
|
||||||
|
for id in &query.ids {
|
||||||
|
separated.push_bind(id.clone());
|
||||||
|
}
|
||||||
|
separated.push_unseparated(")");
|
||||||
|
let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?;
|
||||||
|
rows.iter().map(map_score_row).collect()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl PoolMemberScoreWriteRepository for SqlitePoolMemberScoreRepository {
|
||||||
|
async fn upsert_pool_member_score(
|
||||||
|
&self,
|
||||||
|
score: UpsertPoolMemberScore,
|
||||||
|
) -> Result<StoredPoolMemberScore, DataLayerError> {
|
||||||
|
score.validate()?;
|
||||||
|
let stored = score.into_stored();
|
||||||
|
let score_reason = serde_json::to_string(&stored.score_reason)
|
||||||
|
.map_err(|err| DataLayerError::InvalidInput(err.to_string()))?;
|
||||||
|
sqlx::query(
|
||||||
|
r#"
|
||||||
|
INSERT INTO pool_member_scores (
|
||||||
|
id, pool_kind, pool_id, member_kind, member_id, capability, scope_kind, scope_id,
|
||||||
|
score, hard_state, score_version, score_reason, last_ranked_at, last_scheduled_at,
|
||||||
|
last_success_at, last_failure_at, failure_count, last_probe_attempt_at,
|
||||||
|
last_probe_success_at, last_probe_failure_at, probe_failure_count, probe_status, updated_at
|
||||||
|
) VALUES (
|
||||||
|
?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?
|
||||||
|
)
|
||||||
|
ON CONFLICT(id) DO UPDATE SET
|
||||||
|
pool_kind = excluded.pool_kind,
|
||||||
|
pool_id = excluded.pool_id,
|
||||||
|
member_kind = excluded.member_kind,
|
||||||
|
member_id = excluded.member_id,
|
||||||
|
capability = excluded.capability,
|
||||||
|
scope_kind = excluded.scope_kind,
|
||||||
|
scope_id = excluded.scope_id,
|
||||||
|
score = excluded.score,
|
||||||
|
hard_state = excluded.hard_state,
|
||||||
|
score_version = excluded.score_version,
|
||||||
|
score_reason = excluded.score_reason,
|
||||||
|
last_ranked_at = excluded.last_ranked_at,
|
||||||
|
last_scheduled_at = COALESCE(excluded.last_scheduled_at, pool_member_scores.last_scheduled_at),
|
||||||
|
last_success_at = COALESCE(excluded.last_success_at, pool_member_scores.last_success_at),
|
||||||
|
last_failure_at = COALESCE(excluded.last_failure_at, pool_member_scores.last_failure_at),
|
||||||
|
failure_count = excluded.failure_count,
|
||||||
|
last_probe_attempt_at = COALESCE(excluded.last_probe_attempt_at, pool_member_scores.last_probe_attempt_at),
|
||||||
|
last_probe_success_at = COALESCE(excluded.last_probe_success_at, pool_member_scores.last_probe_success_at),
|
||||||
|
last_probe_failure_at = COALESCE(excluded.last_probe_failure_at, pool_member_scores.last_probe_failure_at),
|
||||||
|
probe_failure_count = excluded.probe_failure_count,
|
||||||
|
probe_status = excluded.probe_status,
|
||||||
|
updated_at = excluded.updated_at
|
||||||
|
"#,
|
||||||
|
)
|
||||||
|
.bind(stored.id.as_str())
|
||||||
|
.bind(stored.pool_kind.as_str())
|
||||||
|
.bind(stored.pool_id.as_str())
|
||||||
|
.bind(stored.member_kind.as_str())
|
||||||
|
.bind(stored.member_id.as_str())
|
||||||
|
.bind(stored.capability.as_str())
|
||||||
|
.bind(stored.scope_kind.as_str())
|
||||||
|
.bind(stored.scope_id.as_deref())
|
||||||
|
.bind(stored.score)
|
||||||
|
.bind(stored.hard_state.as_database())
|
||||||
|
.bind(i64_from_u64(stored.score_version, "pool score version")?)
|
||||||
|
.bind(score_reason)
|
||||||
|
.bind(i64_opt_from_u64(stored.last_ranked_at, "pool score last_ranked_at")?)
|
||||||
|
.bind(i64_opt_from_u64(stored.last_scheduled_at, "pool score last_scheduled_at")?)
|
||||||
|
.bind(i64_opt_from_u64(stored.last_success_at, "pool score last_success_at")?)
|
||||||
|
.bind(i64_opt_from_u64(stored.last_failure_at, "pool score last_failure_at")?)
|
||||||
|
.bind(i64_from_u64(stored.failure_count, "pool score failure_count")?)
|
||||||
|
.bind(i64_opt_from_u64(
|
||||||
|
stored.last_probe_attempt_at,
|
||||||
|
"pool score last_probe_attempt_at",
|
||||||
|
)?)
|
||||||
|
.bind(i64_opt_from_u64(
|
||||||
|
stored.last_probe_success_at,
|
||||||
|
"pool score last_probe_success_at",
|
||||||
|
)?)
|
||||||
|
.bind(i64_opt_from_u64(
|
||||||
|
stored.last_probe_failure_at,
|
||||||
|
"pool score last_probe_failure_at",
|
||||||
|
)?)
|
||||||
|
.bind(i64_from_u64(
|
||||||
|
stored.probe_failure_count,
|
||||||
|
"pool score probe_failure_count",
|
||||||
|
)?)
|
||||||
|
.bind(stored.probe_status.as_database())
|
||||||
|
.bind(i64_from_u64(stored.updated_at, "pool score updated_at")?)
|
||||||
|
.execute(&self.pool)
|
||||||
|
.await
|
||||||
|
.map_sql_err()?;
|
||||||
|
Ok(stored)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn mark_pool_member_probe_in_progress(
|
||||||
|
&self,
|
||||||
|
attempt: PoolMemberProbeAttempt,
|
||||||
|
) -> Result<usize, DataLayerError> {
|
||||||
|
let rows = self
|
||||||
|
.find_scores_by_identity(&attempt.identity, attempt.scope.as_ref())
|
||||||
|
.await?;
|
||||||
|
let count = rows.len();
|
||||||
|
for mut row in rows {
|
||||||
|
row.last_probe_attempt_at = Some(attempt.attempted_at);
|
||||||
|
row.probe_status = PoolMemberProbeStatus::InProgress;
|
||||||
|
row.score_reason =
|
||||||
|
merge_score_reason_patch(row.score_reason, attempt.score_reason_patch.clone());
|
||||||
|
row.updated_at = attempt.attempted_at;
|
||||||
|
self.upsert_pool_member_score(upsert_from_stored(row))
|
||||||
|
.await?;
|
||||||
|
}
|
||||||
|
Ok(count)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn record_pool_member_probe_result(
|
||||||
|
&self,
|
||||||
|
result: PoolMemberProbeResult,
|
||||||
|
) -> Result<usize, DataLayerError> {
|
||||||
|
let rows = self
|
||||||
|
.find_scores_by_identity(&result.identity, result.scope.as_ref())
|
||||||
|
.await?;
|
||||||
|
let count = rows.len();
|
||||||
|
for mut row in rows {
|
||||||
|
row.last_probe_attempt_at = Some(result.attempted_at);
|
||||||
|
row.probe_status = result.probe_status;
|
||||||
|
if result.succeeded {
|
||||||
|
row.last_probe_success_at = Some(result.attempted_at);
|
||||||
|
row.probe_failure_count = 0;
|
||||||
|
} else {
|
||||||
|
row.last_probe_failure_at = Some(result.attempted_at);
|
||||||
|
row.probe_failure_count = row.probe_failure_count.saturating_add(1);
|
||||||
|
}
|
||||||
|
if let Some(hard_state) = result.hard_state {
|
||||||
|
row.hard_state = hard_state;
|
||||||
|
}
|
||||||
|
row.score_reason =
|
||||||
|
merge_score_reason_patch(row.score_reason, result.score_reason_patch.clone());
|
||||||
|
row.updated_at = result.attempted_at;
|
||||||
|
self.upsert_pool_member_score(upsert_from_stored(row))
|
||||||
|
.await?;
|
||||||
|
}
|
||||||
|
Ok(count)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn record_pool_member_schedule_feedback(
|
||||||
|
&self,
|
||||||
|
feedback: PoolMemberScheduleFeedback,
|
||||||
|
) -> Result<usize, DataLayerError> {
|
||||||
|
let rows = self
|
||||||
|
.find_scores_by_identity(&feedback.identity, feedback.scope.as_ref())
|
||||||
|
.await?;
|
||||||
|
let count = rows.len();
|
||||||
|
for mut row in rows {
|
||||||
|
row.last_scheduled_at = Some(feedback.scheduled_at);
|
||||||
|
match feedback.succeeded {
|
||||||
|
Some(true) => row.last_success_at = Some(feedback.scheduled_at),
|
||||||
|
Some(false) => {
|
||||||
|
row.last_failure_at = Some(feedback.scheduled_at);
|
||||||
|
row.failure_count = row.failure_count.saturating_add(1);
|
||||||
|
}
|
||||||
|
None => {}
|
||||||
|
}
|
||||||
|
if let Some(hard_state) = feedback.hard_state {
|
||||||
|
row.hard_state = hard_state;
|
||||||
|
}
|
||||||
|
row.score = score_with_delta(row.score, feedback.score_delta);
|
||||||
|
row.score_reason =
|
||||||
|
merge_score_reason_patch(row.score_reason, feedback.score_reason_patch.clone());
|
||||||
|
row.updated_at = feedback.scheduled_at;
|
||||||
|
self.upsert_pool_member_score(upsert_from_stored(row))
|
||||||
|
.await?;
|
||||||
|
}
|
||||||
|
Ok(count)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn mark_pool_member_hard_state(
|
||||||
|
&self,
|
||||||
|
identity: &PoolMemberIdentity,
|
||||||
|
scope: Option<&PoolScoreScope>,
|
||||||
|
hard_state: PoolMemberHardState,
|
||||||
|
updated_at: u64,
|
||||||
|
) -> Result<usize, DataLayerError> {
|
||||||
|
let rows = self.find_scores_by_identity(identity, scope).await?;
|
||||||
|
let count = rows.len();
|
||||||
|
for mut row in rows {
|
||||||
|
row.hard_state = hard_state;
|
||||||
|
row.updated_at = updated_at;
|
||||||
|
self.upsert_pool_member_score(upsert_from_stored(row))
|
||||||
|
.await?;
|
||||||
|
}
|
||||||
|
Ok(count)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn delete_pool_member_scores_for_member(
|
||||||
|
&self,
|
||||||
|
identity: &PoolMemberIdentity,
|
||||||
|
) -> Result<usize, DataLayerError> {
|
||||||
|
let result = sqlx::query(
|
||||||
|
r#"
|
||||||
|
DELETE FROM pool_member_scores
|
||||||
|
WHERE pool_kind = ? AND pool_id = ? AND member_kind = ? AND member_id = ?
|
||||||
|
"#,
|
||||||
|
)
|
||||||
|
.bind(identity.pool_kind.as_str())
|
||||||
|
.bind(identity.pool_id.as_str())
|
||||||
|
.bind(identity.member_kind.as_str())
|
||||||
|
.bind(identity.member_id.as_str())
|
||||||
|
.execute(&self.pool)
|
||||||
|
.await
|
||||||
|
.map_sql_err()?;
|
||||||
|
Ok(result.rows_affected() as usize)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn map_score_row(row: &SqliteRow) -> Result<StoredPoolMemberScore, DataLayerError> {
|
||||||
|
let score_reason_raw: String = row.try_get("score_reason").map_sql_err()?;
|
||||||
|
Ok(StoredPoolMemberScore {
|
||||||
|
id: row.try_get("id").map_sql_err()?,
|
||||||
|
pool_kind: row.try_get("pool_kind").map_sql_err()?,
|
||||||
|
pool_id: row.try_get("pool_id").map_sql_err()?,
|
||||||
|
member_kind: row.try_get("member_kind").map_sql_err()?,
|
||||||
|
member_id: row.try_get("member_id").map_sql_err()?,
|
||||||
|
capability: row.try_get("capability").map_sql_err()?,
|
||||||
|
scope_kind: row.try_get("scope_kind").map_sql_err()?,
|
||||||
|
scope_id: row.try_get("scope_id").map_sql_err()?,
|
||||||
|
score: row.try_get("score").map_sql_err()?,
|
||||||
|
hard_state: PoolMemberHardState::from_database(
|
||||||
|
row.try_get::<String, _>("hard_state")
|
||||||
|
.map_sql_err()?
|
||||||
|
.as_str(),
|
||||||
|
)?,
|
||||||
|
score_version: u64_from_i64(
|
||||||
|
row.try_get("score_version").map_sql_err()?,
|
||||||
|
"pool_member_scores.score_version",
|
||||||
|
)?,
|
||||||
|
score_reason: serde_json::from_str(&score_reason_raw).unwrap_or(serde_json::Value::Null),
|
||||||
|
last_ranked_at: u64_opt_from_i64(
|
||||||
|
row.try_get("last_ranked_at").map_sql_err()?,
|
||||||
|
"pool_member_scores.last_ranked_at",
|
||||||
|
)?,
|
||||||
|
last_scheduled_at: u64_opt_from_i64(
|
||||||
|
row.try_get("last_scheduled_at").map_sql_err()?,
|
||||||
|
"pool_member_scores.last_scheduled_at",
|
||||||
|
)?,
|
||||||
|
last_success_at: u64_opt_from_i64(
|
||||||
|
row.try_get("last_success_at").map_sql_err()?,
|
||||||
|
"pool_member_scores.last_success_at",
|
||||||
|
)?,
|
||||||
|
last_failure_at: u64_opt_from_i64(
|
||||||
|
row.try_get("last_failure_at").map_sql_err()?,
|
||||||
|
"pool_member_scores.last_failure_at",
|
||||||
|
)?,
|
||||||
|
failure_count: u64_from_i64(
|
||||||
|
row.try_get("failure_count").map_sql_err()?,
|
||||||
|
"pool_member_scores.failure_count",
|
||||||
|
)?,
|
||||||
|
last_probe_attempt_at: u64_opt_from_i64(
|
||||||
|
row.try_get("last_probe_attempt_at").map_sql_err()?,
|
||||||
|
"pool_member_scores.last_probe_attempt_at",
|
||||||
|
)?,
|
||||||
|
last_probe_success_at: u64_opt_from_i64(
|
||||||
|
row.try_get("last_probe_success_at").map_sql_err()?,
|
||||||
|
"pool_member_scores.last_probe_success_at",
|
||||||
|
)?,
|
||||||
|
last_probe_failure_at: u64_opt_from_i64(
|
||||||
|
row.try_get("last_probe_failure_at").map_sql_err()?,
|
||||||
|
"pool_member_scores.last_probe_failure_at",
|
||||||
|
)?,
|
||||||
|
probe_failure_count: u64_from_i64(
|
||||||
|
row.try_get("probe_failure_count").map_sql_err()?,
|
||||||
|
"pool_member_scores.probe_failure_count",
|
||||||
|
)?,
|
||||||
|
probe_status: PoolMemberProbeStatus::from_database(
|
||||||
|
row.try_get::<String, _>("probe_status")
|
||||||
|
.map_sql_err()?
|
||||||
|
.as_str(),
|
||||||
|
)?,
|
||||||
|
updated_at: u64_from_i64(
|
||||||
|
row.try_get("updated_at").map_sql_err()?,
|
||||||
|
"pool_member_scores.updated_at",
|
||||||
|
)?,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn upsert_from_stored(score: StoredPoolMemberScore) -> UpsertPoolMemberScore {
|
||||||
|
UpsertPoolMemberScore {
|
||||||
|
id: score.id,
|
||||||
|
identity: PoolMemberIdentity {
|
||||||
|
pool_kind: score.pool_kind,
|
||||||
|
pool_id: score.pool_id,
|
||||||
|
member_kind: score.member_kind,
|
||||||
|
member_id: score.member_id,
|
||||||
|
},
|
||||||
|
scope: PoolScoreScope {
|
||||||
|
capability: score.capability,
|
||||||
|
scope_kind: score.scope_kind,
|
||||||
|
scope_id: score.scope_id,
|
||||||
|
},
|
||||||
|
score: score.score,
|
||||||
|
hard_state: score.hard_state,
|
||||||
|
score_version: score.score_version,
|
||||||
|
score_reason: score.score_reason,
|
||||||
|
last_ranked_at: score.last_ranked_at,
|
||||||
|
last_scheduled_at: score.last_scheduled_at,
|
||||||
|
last_success_at: score.last_success_at,
|
||||||
|
last_failure_at: score.last_failure_at,
|
||||||
|
failure_count: score.failure_count,
|
||||||
|
last_probe_attempt_at: score.last_probe_attempt_at,
|
||||||
|
last_probe_success_at: score.last_probe_success_at,
|
||||||
|
last_probe_failure_at: score.last_probe_failure_at,
|
||||||
|
probe_failure_count: score.probe_failure_count,
|
||||||
|
probe_status: score.probe_status,
|
||||||
|
updated_at: score.updated_at,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn i64_from_usize(value: usize, field: &str) -> Result<i64, DataLayerError> {
|
||||||
|
i64::try_from(value)
|
||||||
|
.map_err(|_| DataLayerError::InvalidInput(format!("{field} exceeds signed 64-bit range")))
|
||||||
|
}
|
||||||
@@ -192,6 +192,66 @@ export interface PoolKeysPageResponse {
|
|||||||
keys: PoolKeyDetail[]
|
keys: PoolKeyDetail[]
|
||||||
}
|
}
|
||||||
|
|
||||||
|
export type PoolScoreHardState =
|
||||||
|
| 'available'
|
||||||
|
| 'unknown'
|
||||||
|
| 'cooldown'
|
||||||
|
| 'quota_exhausted'
|
||||||
|
| 'auth_invalid'
|
||||||
|
| 'banned'
|
||||||
|
| 'inactive'
|
||||||
|
|
||||||
|
export type PoolScoreProbeStatus = 'never' | 'ok' | 'failed' | 'stale' | 'in_progress'
|
||||||
|
|
||||||
|
export interface PoolScoreKeySummary {
|
||||||
|
id: string
|
||||||
|
name: string
|
||||||
|
auth_type: string
|
||||||
|
is_active: boolean
|
||||||
|
internal_priority: number
|
||||||
|
last_used_at: number | null
|
||||||
|
}
|
||||||
|
|
||||||
|
export interface PoolMemberScoreItem {
|
||||||
|
id: string
|
||||||
|
pool_kind: string
|
||||||
|
pool_id: string
|
||||||
|
member_kind: string
|
||||||
|
member_id: string
|
||||||
|
capability: string
|
||||||
|
scope_kind: string
|
||||||
|
scope_id: string | null
|
||||||
|
score: number
|
||||||
|
hard_state: PoolScoreHardState
|
||||||
|
score_version: number
|
||||||
|
score_reason: Record<string, unknown> | null
|
||||||
|
last_ranked_at: number | null
|
||||||
|
last_scheduled_at: number | null
|
||||||
|
last_success_at: number | null
|
||||||
|
last_failure_at: number | null
|
||||||
|
failure_count: number
|
||||||
|
last_probe_attempt_at: number | null
|
||||||
|
last_probe_success_at: number | null
|
||||||
|
last_probe_failure_at: number | null
|
||||||
|
probe_failure_count: number
|
||||||
|
probe_status: PoolScoreProbeStatus
|
||||||
|
updated_at: number
|
||||||
|
key?: PoolScoreKeySummary | null
|
||||||
|
}
|
||||||
|
|
||||||
|
export interface PoolScoresResponse {
|
||||||
|
provider_id: string
|
||||||
|
page: number
|
||||||
|
page_size: number
|
||||||
|
filters: {
|
||||||
|
api_format?: string | null
|
||||||
|
model_id?: string | null
|
||||||
|
hard_state?: string | null
|
||||||
|
probe_status?: string | null
|
||||||
|
}
|
||||||
|
items: PoolMemberScoreItem[]
|
||||||
|
}
|
||||||
|
|
||||||
export interface PoolKeysQuery {
|
export interface PoolKeysQuery {
|
||||||
page?: number
|
page?: number
|
||||||
page_size?: number
|
page_size?: number
|
||||||
@@ -203,6 +263,15 @@ export interface PoolKeysQuery {
|
|||||||
sort_order?: 'asc' | 'desc'
|
sort_order?: 'asc' | 'desc'
|
||||||
}
|
}
|
||||||
|
|
||||||
|
export interface PoolScoresQuery {
|
||||||
|
page?: number
|
||||||
|
page_size?: number
|
||||||
|
api_format?: string
|
||||||
|
model_id?: string
|
||||||
|
hard_state?: string
|
||||||
|
probe_status?: string
|
||||||
|
}
|
||||||
|
|
||||||
export interface PoolKeySelectionRequest {
|
export interface PoolKeySelectionRequest {
|
||||||
search?: string
|
search?: string
|
||||||
quick_selectors?: string[]
|
quick_selectors?: string[]
|
||||||
@@ -293,6 +362,29 @@ export async function listPoolKeys(
|
|||||||
)
|
)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
export async function listPoolScores(
|
||||||
|
providerId: string,
|
||||||
|
params: PoolScoresQuery = {},
|
||||||
|
options: PoolReadOptions = {},
|
||||||
|
): Promise<PoolScoresResponse> {
|
||||||
|
const normalizedParams = { ...params }
|
||||||
|
const cacheKey = buildCacheKey(
|
||||||
|
`pool:scores:${providerId}`,
|
||||||
|
normalizedParams as Record<string, unknown>,
|
||||||
|
)
|
||||||
|
return cachedRequest(
|
||||||
|
cacheKey,
|
||||||
|
async () => {
|
||||||
|
const response = await client.get<PoolScoresResponse>(
|
||||||
|
`/api/admin/pool/${providerId}/scores`,
|
||||||
|
{ params: normalizedParams },
|
||||||
|
)
|
||||||
|
return response.data
|
||||||
|
},
|
||||||
|
options.cacheTtlMs ?? 0,
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
export async function resolvePoolKeySelection(
|
export async function resolvePoolKeySelection(
|
||||||
providerId: string,
|
providerId: string,
|
||||||
body: PoolKeySelectionRequest,
|
body: PoolKeySelectionRequest,
|
||||||
|
|||||||
@@ -521,6 +521,24 @@ export interface SchedulingPresetItem {
|
|||||||
mode?: string | null
|
mode?: string | null
|
||||||
}
|
}
|
||||||
|
|
||||||
|
export interface PoolScoreWeights {
|
||||||
|
manual_priority?: number | null
|
||||||
|
health?: number | null
|
||||||
|
probe_freshness?: number | null
|
||||||
|
quota_remaining?: number | null
|
||||||
|
latency?: number | null
|
||||||
|
cost_lru?: number | null
|
||||||
|
}
|
||||||
|
|
||||||
|
export interface PoolScoreRules {
|
||||||
|
weights?: PoolScoreWeights | null
|
||||||
|
probe_freshness_ttl_seconds?: number | null
|
||||||
|
unschedulable_score_cap?: number | null
|
||||||
|
probe_failure_penalty?: number | null
|
||||||
|
request_failure_penalty?: number | null
|
||||||
|
probe_failure_cooldown_threshold?: number | null
|
||||||
|
}
|
||||||
|
|
||||||
export interface PoolAdvancedConfig {
|
export interface PoolAdvancedConfig {
|
||||||
global_priority?: number | null
|
global_priority?: number | null
|
||||||
sticky_session_ttl_seconds?: number | null
|
sticky_session_ttl_seconds?: number | null
|
||||||
@@ -548,6 +566,10 @@ export interface PoolAdvancedConfig {
|
|||||||
health_policy_enabled?: boolean
|
health_policy_enabled?: boolean
|
||||||
unschedulable_rules?: Array<Record<string, unknown>> | null
|
unschedulable_rules?: Array<Record<string, unknown>> | null
|
||||||
batch_concurrency?: number | null
|
batch_concurrency?: number | null
|
||||||
|
probe_concurrency?: number | null
|
||||||
|
score_top_n?: number | null
|
||||||
|
score_fallback_scan_limit?: number | null
|
||||||
|
score_rules?: PoolScoreRules | null
|
||||||
probing_enabled?: boolean
|
probing_enabled?: boolean
|
||||||
probing_interval_minutes?: number | null
|
probing_interval_minutes?: number | null
|
||||||
auto_remove_banned_keys?: boolean
|
auto_remove_banned_keys?: boolean
|
||||||
|
|||||||
@@ -230,7 +230,7 @@
|
|||||||
</p>
|
</p>
|
||||||
</div>
|
</div>
|
||||||
|
|
||||||
<div class="rounded-xl bg-muted/30 p-4">
|
<div class="grid gap-3 rounded-xl bg-muted/30 p-4 sm:grid-cols-2">
|
||||||
<div class="space-y-1.5">
|
<div class="space-y-1.5">
|
||||||
<Label>
|
<Label>
|
||||||
并发数
|
并发数
|
||||||
@@ -247,10 +247,204 @@
|
|||||||
为空时沿用默认值;数值越大,批量操作越快,但会增加瞬时请求压力。
|
为空时沿用默认值;数值越大,批量操作越快,但会增加瞬时请求压力。
|
||||||
</p>
|
</p>
|
||||||
</div>
|
</div>
|
||||||
|
<div class="space-y-1.5">
|
||||||
|
<Label>
|
||||||
|
探测并发
|
||||||
|
</Label>
|
||||||
|
<Input
|
||||||
|
:model-value="form.probe_concurrency ?? ''"
|
||||||
|
type="number"
|
||||||
|
min="1"
|
||||||
|
max="64"
|
||||||
|
placeholder="4"
|
||||||
|
@update:model-value="(v) => form.probe_concurrency = parseNum(v)"
|
||||||
|
/>
|
||||||
|
</div>
|
||||||
|
<div class="space-y-1.5">
|
||||||
|
<Label>
|
||||||
|
评分 Top-N
|
||||||
|
</Label>
|
||||||
|
<Input
|
||||||
|
:model-value="form.score_top_n ?? ''"
|
||||||
|
type="number"
|
||||||
|
min="1"
|
||||||
|
max="4096"
|
||||||
|
placeholder="128"
|
||||||
|
@update:model-value="(v) => form.score_top_n = parseNum(v)"
|
||||||
|
/>
|
||||||
|
</div>
|
||||||
|
<div class="space-y-1.5">
|
||||||
|
<Label>
|
||||||
|
回退扫描
|
||||||
|
</Label>
|
||||||
|
<Input
|
||||||
|
:model-value="form.score_fallback_scan_limit ?? ''"
|
||||||
|
type="number"
|
||||||
|
min="1"
|
||||||
|
max="100000"
|
||||||
|
placeholder="1024"
|
||||||
|
@update:model-value="(v) => form.score_fallback_scan_limit = parseNum(v)"
|
||||||
|
/>
|
||||||
|
</div>
|
||||||
</div>
|
</div>
|
||||||
</section>
|
</section>
|
||||||
</div>
|
</div>
|
||||||
|
|
||||||
|
<section class="space-y-4 rounded-2xl border border-border/60 bg-card/70 p-4 sm:p-5">
|
||||||
|
<div class="space-y-1">
|
||||||
|
<div class="flex flex-wrap items-center gap-2">
|
||||||
|
<h3 class="text-sm font-semibold">
|
||||||
|
分数规则
|
||||||
|
</h3>
|
||||||
|
<span class="rounded-full bg-muted px-2 py-0.5 text-[11px] text-muted-foreground">
|
||||||
|
候选排序
|
||||||
|
</span>
|
||||||
|
</div>
|
||||||
|
<p class="text-xs leading-5 text-muted-foreground">
|
||||||
|
调整主动探测、健康、额度、延迟和使用成本进入号池候选排序时的权重。
|
||||||
|
</p>
|
||||||
|
</div>
|
||||||
|
|
||||||
|
<div class="grid gap-3 lg:grid-cols-3">
|
||||||
|
<div class="space-y-1.5">
|
||||||
|
<Label>优先级权重</Label>
|
||||||
|
<Input
|
||||||
|
:model-value="form.score_weight_manual_priority ?? ''"
|
||||||
|
type="number"
|
||||||
|
min="0"
|
||||||
|
max="1"
|
||||||
|
step="0.01"
|
||||||
|
placeholder="0.30"
|
||||||
|
@update:model-value="(v) => form.score_weight_manual_priority = parseNum(v)"
|
||||||
|
/>
|
||||||
|
</div>
|
||||||
|
<div class="space-y-1.5">
|
||||||
|
<Label>健康权重</Label>
|
||||||
|
<Input
|
||||||
|
:model-value="form.score_weight_health ?? ''"
|
||||||
|
type="number"
|
||||||
|
min="0"
|
||||||
|
max="1"
|
||||||
|
step="0.01"
|
||||||
|
placeholder="0.20"
|
||||||
|
@update:model-value="(v) => form.score_weight_health = parseNum(v)"
|
||||||
|
/>
|
||||||
|
</div>
|
||||||
|
<div class="space-y-1.5">
|
||||||
|
<Label>探测新鲜度权重</Label>
|
||||||
|
<Input
|
||||||
|
:model-value="form.score_weight_probe_freshness ?? ''"
|
||||||
|
type="number"
|
||||||
|
min="0"
|
||||||
|
max="1"
|
||||||
|
step="0.01"
|
||||||
|
placeholder="0.15"
|
||||||
|
@update:model-value="(v) => form.score_weight_probe_freshness = parseNum(v)"
|
||||||
|
/>
|
||||||
|
</div>
|
||||||
|
<div class="space-y-1.5">
|
||||||
|
<Label>额度剩余权重</Label>
|
||||||
|
<Input
|
||||||
|
:model-value="form.score_weight_quota_remaining ?? ''"
|
||||||
|
type="number"
|
||||||
|
min="0"
|
||||||
|
max="1"
|
||||||
|
step="0.01"
|
||||||
|
placeholder="0.15"
|
||||||
|
@update:model-value="(v) => form.score_weight_quota_remaining = parseNum(v)"
|
||||||
|
/>
|
||||||
|
</div>
|
||||||
|
<div class="space-y-1.5">
|
||||||
|
<Label>延迟权重</Label>
|
||||||
|
<Input
|
||||||
|
:model-value="form.score_weight_latency ?? ''"
|
||||||
|
type="number"
|
||||||
|
min="0"
|
||||||
|
max="1"
|
||||||
|
step="0.01"
|
||||||
|
placeholder="0.10"
|
||||||
|
@update:model-value="(v) => form.score_weight_latency = parseNum(v)"
|
||||||
|
/>
|
||||||
|
</div>
|
||||||
|
<div class="space-y-1.5">
|
||||||
|
<Label>成本/LRU 权重</Label>
|
||||||
|
<Input
|
||||||
|
:model-value="form.score_weight_cost_lru ?? ''"
|
||||||
|
type="number"
|
||||||
|
min="0"
|
||||||
|
max="1"
|
||||||
|
step="0.01"
|
||||||
|
placeholder="0.10"
|
||||||
|
@update:model-value="(v) => form.score_weight_cost_lru = parseNum(v)"
|
||||||
|
/>
|
||||||
|
</div>
|
||||||
|
</div>
|
||||||
|
|
||||||
|
<div class="grid gap-3 sm:grid-cols-2 lg:grid-cols-3">
|
||||||
|
<div class="space-y-1.5">
|
||||||
|
<Label>
|
||||||
|
探测新鲜度 TTL
|
||||||
|
<span class="text-xs text-muted-foreground">(秒)</span>
|
||||||
|
</Label>
|
||||||
|
<Input
|
||||||
|
:model-value="form.probe_freshness_ttl_seconds ?? ''"
|
||||||
|
type="number"
|
||||||
|
min="1"
|
||||||
|
max="604800"
|
||||||
|
placeholder="1800"
|
||||||
|
@update:model-value="(v) => form.probe_freshness_ttl_seconds = parseNum(v)"
|
||||||
|
/>
|
||||||
|
</div>
|
||||||
|
<div class="space-y-1.5">
|
||||||
|
<Label>探测失败惩罚</Label>
|
||||||
|
<Input
|
||||||
|
:model-value="form.probe_failure_penalty ?? ''"
|
||||||
|
type="number"
|
||||||
|
min="0"
|
||||||
|
max="1"
|
||||||
|
step="0.01"
|
||||||
|
placeholder="0.05"
|
||||||
|
@update:model-value="(v) => form.probe_failure_penalty = parseNum(v)"
|
||||||
|
/>
|
||||||
|
</div>
|
||||||
|
<div class="space-y-1.5">
|
||||||
|
<Label>请求失败惩罚</Label>
|
||||||
|
<Input
|
||||||
|
:model-value="form.request_failure_penalty ?? ''"
|
||||||
|
type="number"
|
||||||
|
min="0"
|
||||||
|
max="1"
|
||||||
|
step="0.001"
|
||||||
|
placeholder="0.005"
|
||||||
|
@update:model-value="(v) => form.request_failure_penalty = parseNum(v)"
|
||||||
|
/>
|
||||||
|
</div>
|
||||||
|
<div class="space-y-1.5">
|
||||||
|
<Label>探测失败冷却阈值</Label>
|
||||||
|
<Input
|
||||||
|
:model-value="form.probe_failure_cooldown_threshold ?? ''"
|
||||||
|
type="number"
|
||||||
|
min="0"
|
||||||
|
max="100"
|
||||||
|
placeholder="3"
|
||||||
|
@update:model-value="(v) => form.probe_failure_cooldown_threshold = parseNum(v)"
|
||||||
|
/>
|
||||||
|
</div>
|
||||||
|
<div class="space-y-1.5">
|
||||||
|
<Label>不可调度分数上限</Label>
|
||||||
|
<Input
|
||||||
|
:model-value="form.unschedulable_score_cap ?? ''"
|
||||||
|
type="number"
|
||||||
|
min="0"
|
||||||
|
max="1"
|
||||||
|
step="0.01"
|
||||||
|
placeholder="0.05"
|
||||||
|
@update:model-value="(v) => form.unschedulable_score_cap = parseNum(v)"
|
||||||
|
/>
|
||||||
|
</div>
|
||||||
|
</div>
|
||||||
|
</section>
|
||||||
|
|
||||||
<section
|
<section
|
||||||
v-if="isClaudeCode"
|
v-if="isClaudeCode"
|
||||||
class="space-y-4 rounded-2xl border border-border/60 bg-card/70 p-4 sm:p-5"
|
class="space-y-4 rounded-2xl border border-border/60 bg-card/70 p-4 sm:p-5"
|
||||||
@@ -464,6 +658,20 @@ const form = ref({
|
|||||||
cost_limit_per_key_tokens: null as number | null | undefined,
|
cost_limit_per_key_tokens: null as number | null | undefined,
|
||||||
cost_soft_threshold_percent: null as number | null | undefined,
|
cost_soft_threshold_percent: null as number | null | undefined,
|
||||||
batch_concurrency: null as number | null | undefined,
|
batch_concurrency: null as number | null | undefined,
|
||||||
|
probe_concurrency: null as number | null | undefined,
|
||||||
|
score_top_n: null as number | null | undefined,
|
||||||
|
score_fallback_scan_limit: null as number | null | undefined,
|
||||||
|
score_weight_manual_priority: null as number | null | undefined,
|
||||||
|
score_weight_health: null as number | null | undefined,
|
||||||
|
score_weight_probe_freshness: null as number | null | undefined,
|
||||||
|
score_weight_quota_remaining: null as number | null | undefined,
|
||||||
|
score_weight_latency: null as number | null | undefined,
|
||||||
|
score_weight_cost_lru: null as number | null | undefined,
|
||||||
|
probe_freshness_ttl_seconds: null as number | null | undefined,
|
||||||
|
unschedulable_score_cap: null as number | null | undefined,
|
||||||
|
probe_failure_penalty: null as number | null | undefined,
|
||||||
|
request_failure_penalty: null as number | null | undefined,
|
||||||
|
probe_failure_cooldown_threshold: null as number | null | undefined,
|
||||||
probing_enabled: false,
|
probing_enabled: false,
|
||||||
probing_interval_minutes: null as number | null | undefined,
|
probing_interval_minutes: null as number | null | undefined,
|
||||||
auto_remove_banned_keys: false,
|
auto_remove_banned_keys: false,
|
||||||
@@ -529,6 +737,8 @@ watch(() => props.modelValue, (open) => {
|
|||||||
if (!open) return
|
if (!open) return
|
||||||
|
|
||||||
const cfg = props.currentConfig
|
const cfg = props.currentConfig
|
||||||
|
const scoreRules = cfg?.score_rules
|
||||||
|
const scoreWeights = scoreRules?.weights
|
||||||
form.value = {
|
form.value = {
|
||||||
global_priority: cfg?.global_priority ?? null,
|
global_priority: cfg?.global_priority ?? null,
|
||||||
sticky_session_ttl_seconds: cfg?.sticky_session_ttl_seconds ?? null,
|
sticky_session_ttl_seconds: cfg?.sticky_session_ttl_seconds ?? null,
|
||||||
@@ -539,6 +749,20 @@ watch(() => props.modelValue, (open) => {
|
|||||||
cost_limit_per_key_tokens: cfg?.cost_limit_per_key_tokens ?? null,
|
cost_limit_per_key_tokens: cfg?.cost_limit_per_key_tokens ?? null,
|
||||||
cost_soft_threshold_percent: cfg?.cost_soft_threshold_percent ?? null,
|
cost_soft_threshold_percent: cfg?.cost_soft_threshold_percent ?? null,
|
||||||
batch_concurrency: cfg?.batch_concurrency ?? null,
|
batch_concurrency: cfg?.batch_concurrency ?? null,
|
||||||
|
probe_concurrency: cfg?.probe_concurrency ?? null,
|
||||||
|
score_top_n: cfg?.score_top_n ?? null,
|
||||||
|
score_fallback_scan_limit: cfg?.score_fallback_scan_limit ?? null,
|
||||||
|
score_weight_manual_priority: scoreWeights?.manual_priority ?? null,
|
||||||
|
score_weight_health: scoreWeights?.health ?? null,
|
||||||
|
score_weight_probe_freshness: scoreWeights?.probe_freshness ?? null,
|
||||||
|
score_weight_quota_remaining: scoreWeights?.quota_remaining ?? null,
|
||||||
|
score_weight_latency: scoreWeights?.latency ?? null,
|
||||||
|
score_weight_cost_lru: scoreWeights?.cost_lru ?? null,
|
||||||
|
probe_freshness_ttl_seconds: scoreRules?.probe_freshness_ttl_seconds ?? null,
|
||||||
|
unschedulable_score_cap: scoreRules?.unschedulable_score_cap ?? null,
|
||||||
|
probe_failure_penalty: scoreRules?.probe_failure_penalty ?? null,
|
||||||
|
request_failure_penalty: scoreRules?.request_failure_penalty ?? null,
|
||||||
|
probe_failure_cooldown_threshold: scoreRules?.probe_failure_cooldown_threshold ?? null,
|
||||||
probing_enabled: cfg?.probing_enabled ?? false,
|
probing_enabled: cfg?.probing_enabled ?? false,
|
||||||
probing_interval_minutes: cfg?.probing_interval_minutes ?? null,
|
probing_interval_minutes: cfg?.probing_interval_minutes ?? null,
|
||||||
auto_remove_banned_keys: cfg?.auto_remove_banned_keys ?? false,
|
auto_remove_banned_keys: cfg?.auto_remove_banned_keys ?? false,
|
||||||
@@ -560,6 +784,23 @@ watch(() => props.modelValue, (open) => {
|
|||||||
async function handleSave() {
|
async function handleSave() {
|
||||||
loading.value = true
|
loading.value = true
|
||||||
try {
|
try {
|
||||||
|
const scoreRules = {
|
||||||
|
...(props.currentConfig?.score_rules ?? {}),
|
||||||
|
weights: {
|
||||||
|
...(props.currentConfig?.score_rules?.weights ?? {}),
|
||||||
|
manual_priority: form.value.score_weight_manual_priority ?? undefined,
|
||||||
|
health: form.value.score_weight_health ?? undefined,
|
||||||
|
probe_freshness: form.value.score_weight_probe_freshness ?? undefined,
|
||||||
|
quota_remaining: form.value.score_weight_quota_remaining ?? undefined,
|
||||||
|
latency: form.value.score_weight_latency ?? undefined,
|
||||||
|
cost_lru: form.value.score_weight_cost_lru ?? undefined,
|
||||||
|
},
|
||||||
|
probe_freshness_ttl_seconds: form.value.probe_freshness_ttl_seconds ?? undefined,
|
||||||
|
unschedulable_score_cap: form.value.unschedulable_score_cap ?? undefined,
|
||||||
|
probe_failure_penalty: form.value.probe_failure_penalty ?? undefined,
|
||||||
|
request_failure_penalty: form.value.request_failure_penalty ?? undefined,
|
||||||
|
probe_failure_cooldown_threshold: form.value.probe_failure_cooldown_threshold ?? undefined,
|
||||||
|
}
|
||||||
// 合并已有配置(保留 scheduling_presets 等不在此对话框编辑的字段)
|
// 合并已有配置(保留 scheduling_presets 等不在此对话框编辑的字段)
|
||||||
const poolAdvanced: Record<string, unknown> = {
|
const poolAdvanced: Record<string, unknown> = {
|
||||||
...(props.currentConfig ?? {}),
|
...(props.currentConfig ?? {}),
|
||||||
@@ -572,6 +813,10 @@ async function handleSave() {
|
|||||||
overload_cooldown_seconds: form.value.overload_cooldown_seconds ?? undefined,
|
overload_cooldown_seconds: form.value.overload_cooldown_seconds ?? undefined,
|
||||||
health_policy_enabled: form.value.health_policy_enabled,
|
health_policy_enabled: form.value.health_policy_enabled,
|
||||||
batch_concurrency: form.value.batch_concurrency ?? undefined,
|
batch_concurrency: form.value.batch_concurrency ?? undefined,
|
||||||
|
probe_concurrency: form.value.probe_concurrency ?? undefined,
|
||||||
|
score_top_n: form.value.score_top_n ?? undefined,
|
||||||
|
score_fallback_scan_limit: form.value.score_fallback_scan_limit ?? undefined,
|
||||||
|
score_rules: scoreRules,
|
||||||
probing_enabled: form.value.probing_enabled,
|
probing_enabled: form.value.probing_enabled,
|
||||||
probing_interval_minutes: form.value.probing_enabled
|
probing_interval_minutes: form.value.probing_enabled
|
||||||
? (form.value.probing_interval_minutes ?? undefined)
|
? (form.value.probing_interval_minutes ?? undefined)
|
||||||
|
|||||||
@@ -1189,6 +1189,170 @@
|
|||||||
</template>
|
</template>
|
||||||
</Card>
|
</Card>
|
||||||
|
|
||||||
|
<Card
|
||||||
|
v-if="selectedProviderId"
|
||||||
|
variant="default"
|
||||||
|
class="overflow-hidden"
|
||||||
|
>
|
||||||
|
<div class="flex flex-col gap-3 border-b border-border/60 px-4 py-3 sm:flex-row sm:items-center sm:justify-between sm:px-6 sm:py-3.5">
|
||||||
|
<div class="min-w-0">
|
||||||
|
<h3 class="text-base font-semibold">
|
||||||
|
评分排序
|
||||||
|
</h3>
|
||||||
|
<p class="mt-1 text-xs text-muted-foreground">
|
||||||
|
Top {{ poolScorePage.items.length }} / page {{ poolScorePage.page || 1 }}
|
||||||
|
</p>
|
||||||
|
</div>
|
||||||
|
<RefreshButton
|
||||||
|
:loading="scoresLoading"
|
||||||
|
title="刷新评分"
|
||||||
|
@click="() => loadScores()"
|
||||||
|
/>
|
||||||
|
</div>
|
||||||
|
|
||||||
|
<div class="grid gap-2 border-b border-border/60 bg-muted/20 px-4 py-3 sm:grid-cols-4 sm:px-6">
|
||||||
|
<Input
|
||||||
|
v-model="scoreApiFormatFilter"
|
||||||
|
class="h-8 text-xs"
|
||||||
|
placeholder="api_format"
|
||||||
|
/>
|
||||||
|
<Input
|
||||||
|
v-model="scoreModelIdFilter"
|
||||||
|
class="h-8 text-xs"
|
||||||
|
placeholder="model_id"
|
||||||
|
/>
|
||||||
|
<Select v-model="scoreHardStateFilter">
|
||||||
|
<SelectTrigger class="h-8 text-xs border-border/60">
|
||||||
|
<SelectValue placeholder="硬状态" />
|
||||||
|
</SelectTrigger>
|
||||||
|
<SelectContent>
|
||||||
|
<SelectItem
|
||||||
|
v-for="item in poolScoreHardStateOptions"
|
||||||
|
:key="item.value"
|
||||||
|
:value="item.value"
|
||||||
|
>
|
||||||
|
{{ item.label }}
|
||||||
|
</SelectItem>
|
||||||
|
</SelectContent>
|
||||||
|
</Select>
|
||||||
|
<Select v-model="scoreProbeStatusFilter">
|
||||||
|
<SelectTrigger class="h-8 text-xs border-border/60">
|
||||||
|
<SelectValue placeholder="探测" />
|
||||||
|
</SelectTrigger>
|
||||||
|
<SelectContent>
|
||||||
|
<SelectItem
|
||||||
|
v-for="item in poolScoreProbeStatusOptions"
|
||||||
|
:key="item.value"
|
||||||
|
:value="item.value"
|
||||||
|
>
|
||||||
|
{{ item.label }}
|
||||||
|
</SelectItem>
|
||||||
|
</SelectContent>
|
||||||
|
</Select>
|
||||||
|
</div>
|
||||||
|
|
||||||
|
<div
|
||||||
|
v-if="scoresLoading && poolScorePage.items.length === 0"
|
||||||
|
class="flex items-center justify-center py-10"
|
||||||
|
>
|
||||||
|
<div class="animate-spin rounded-full h-7 w-7 border-b-2 border-primary" />
|
||||||
|
</div>
|
||||||
|
<div
|
||||||
|
v-else-if="poolScorePage.items.length === 0"
|
||||||
|
class="flex flex-col items-center justify-center py-10 text-center"
|
||||||
|
>
|
||||||
|
<div class="mx-auto flex h-12 w-12 items-center justify-center rounded-full bg-muted">
|
||||||
|
<SlidersHorizontal class="h-6 w-6 text-muted-foreground" />
|
||||||
|
</div>
|
||||||
|
<p class="mt-3 text-sm text-muted-foreground">
|
||||||
|
{{ hasPoolScoreFilters ? '未找到匹配评分' : '暂无评分数据' }}
|
||||||
|
</p>
|
||||||
|
</div>
|
||||||
|
<div
|
||||||
|
v-else
|
||||||
|
class="overflow-x-auto"
|
||||||
|
>
|
||||||
|
<Table class="w-full min-w-[760px] table-fixed">
|
||||||
|
<TableHeader>
|
||||||
|
<TableRow class="border-b border-border/60 hover:bg-transparent">
|
||||||
|
<TableHead class="w-[24%] font-semibold">
|
||||||
|
账号
|
||||||
|
</TableHead>
|
||||||
|
<TableHead class="w-[24%] font-semibold">
|
||||||
|
范围
|
||||||
|
</TableHead>
|
||||||
|
<TableHead class="w-[12%] text-center font-semibold">
|
||||||
|
分数
|
||||||
|
</TableHead>
|
||||||
|
<TableHead class="w-[12%] text-center font-semibold">
|
||||||
|
状态
|
||||||
|
</TableHead>
|
||||||
|
<TableHead class="w-[12%] text-center font-semibold">
|
||||||
|
探测
|
||||||
|
</TableHead>
|
||||||
|
<TableHead class="w-[16%] text-center font-semibold">
|
||||||
|
更新时间
|
||||||
|
</TableHead>
|
||||||
|
</TableRow>
|
||||||
|
</TableHeader>
|
||||||
|
<TableBody>
|
||||||
|
<TableRow
|
||||||
|
v-for="score in poolScorePage.items"
|
||||||
|
:key="score.id"
|
||||||
|
class="border-b border-border/40 last:border-b-0 hover:bg-muted/30"
|
||||||
|
>
|
||||||
|
<TableCell class="py-3">
|
||||||
|
<div class="min-w-0">
|
||||||
|
<div class="truncate text-sm">
|
||||||
|
{{ score.key?.name || score.member_id }}
|
||||||
|
</div>
|
||||||
|
<div class="truncate font-mono text-[10px] text-muted-foreground">
|
||||||
|
{{ score.member_id }}
|
||||||
|
</div>
|
||||||
|
</div>
|
||||||
|
</TableCell>
|
||||||
|
<TableCell class="py-3">
|
||||||
|
<div class="min-w-0 text-xs">
|
||||||
|
<div class="truncate text-foreground/90">
|
||||||
|
{{ score.capability || '-' }}
|
||||||
|
</div>
|
||||||
|
<div class="truncate text-[10px] text-muted-foreground">
|
||||||
|
{{ score.scope_kind }}: {{ score.scope_id || '*' }}
|
||||||
|
</div>
|
||||||
|
</div>
|
||||||
|
</TableCell>
|
||||||
|
<TableCell class="py-3 text-center">
|
||||||
|
<span class="font-mono text-xs tabular-nums">
|
||||||
|
{{ formatPoolScore(score.score) }}
|
||||||
|
</span>
|
||||||
|
</TableCell>
|
||||||
|
<TableCell class="py-3 text-center">
|
||||||
|
<Badge
|
||||||
|
variant="outline"
|
||||||
|
class="text-[10px]"
|
||||||
|
>
|
||||||
|
{{ getPoolScoreHardStateLabel(score.hard_state) }}
|
||||||
|
</Badge>
|
||||||
|
</TableCell>
|
||||||
|
<TableCell class="py-3 text-center">
|
||||||
|
<Badge
|
||||||
|
variant="secondary"
|
||||||
|
class="text-[10px]"
|
||||||
|
>
|
||||||
|
{{ getPoolScoreProbeStatusLabel(score.probe_status) }}
|
||||||
|
</Badge>
|
||||||
|
</TableCell>
|
||||||
|
<TableCell class="py-3 text-center">
|
||||||
|
<span class="text-[10px] text-muted-foreground">
|
||||||
|
{{ formatUnixSeconds(score.updated_at) }}
|
||||||
|
</span>
|
||||||
|
</TableCell>
|
||||||
|
</TableRow>
|
||||||
|
</TableBody>
|
||||||
|
</Table>
|
||||||
|
</div>
|
||||||
|
</Card>
|
||||||
|
|
||||||
<!-- Dialogs -->
|
<!-- Dialogs -->
|
||||||
<OAuthAccountDialog
|
<OAuthAccountDialog
|
||||||
v-if="selectedProviderId"
|
v-if="selectedProviderId"
|
||||||
@@ -1309,6 +1473,7 @@ import {
|
|||||||
getPoolOverview,
|
getPoolOverview,
|
||||||
getPoolSchedulingPresets,
|
getPoolSchedulingPresets,
|
||||||
listPoolKeys,
|
listPoolKeys,
|
||||||
|
listPoolScores,
|
||||||
clearPoolCooldown,
|
clearPoolCooldown,
|
||||||
} from '@/api/endpoints/pool'
|
} from '@/api/endpoints/pool'
|
||||||
import {
|
import {
|
||||||
@@ -1324,7 +1489,9 @@ import type {
|
|||||||
PoolOverviewItem,
|
PoolOverviewItem,
|
||||||
PoolKeyDetail,
|
PoolKeyDetail,
|
||||||
PoolKeysPageResponse,
|
PoolKeysPageResponse,
|
||||||
|
PoolMemberScoreItem,
|
||||||
PoolPresetMeta,
|
PoolPresetMeta,
|
||||||
|
PoolScoresResponse,
|
||||||
} from '@/api/endpoints/pool'
|
} from '@/api/endpoints/pool'
|
||||||
import type {
|
import type {
|
||||||
ClaudeCodeAdvancedConfig,
|
ClaudeCodeAdvancedConfig,
|
||||||
@@ -1420,11 +1587,14 @@ let overviewRequestId = 0
|
|||||||
let selectProviderRequestId = 0
|
let selectProviderRequestId = 0
|
||||||
let providerDataRequestId = 0
|
let providerDataRequestId = 0
|
||||||
let keysRequestId = 0
|
let keysRequestId = 0
|
||||||
|
let scoresRequestId = 0
|
||||||
let keysSearchDebounceTimer: number | null = null
|
let keysSearchDebounceTimer: number | null = null
|
||||||
|
let scoresFilterDebounceTimer: number | null = null
|
||||||
let suppressFiltersWatch = false
|
let suppressFiltersWatch = false
|
||||||
let hasHydratedInitialProviderSelection = false
|
let hasHydratedInitialProviderSelection = false
|
||||||
const POOL_OVERVIEW_CACHE_TTL_MS = 10 * 1000
|
const POOL_OVERVIEW_CACHE_TTL_MS = 10 * 1000
|
||||||
const POOL_KEYS_CACHE_TTL_MS = 10 * 1000
|
const POOL_KEYS_CACHE_TTL_MS = 10 * 1000
|
||||||
|
const POOL_SCORES_CACHE_TTL_MS = 10 * 1000
|
||||||
const POOL_SCHEDULING_PRESETS_CACHE_TTL_MS = 5 * 60 * 1000
|
const POOL_SCHEDULING_PRESETS_CACHE_TTL_MS = 5 * 60 * 1000
|
||||||
const poolKeyStatusFilterOptions: Array<{ value: PoolManagementViewState['status'], label: string }> = [
|
const poolKeyStatusFilterOptions: Array<{ value: PoolManagementViewState['status'], label: string }> = [
|
||||||
{ value: 'all', label: '全部状态' },
|
{ value: 'all', label: '全部状态' },
|
||||||
@@ -1432,6 +1602,24 @@ const poolKeyStatusFilterOptions: Array<{ value: PoolManagementViewState['status
|
|||||||
{ value: 'cooldown', label: '冷却中' },
|
{ value: 'cooldown', label: '冷却中' },
|
||||||
{ value: 'inactive', label: '禁用' },
|
{ value: 'inactive', label: '禁用' },
|
||||||
]
|
]
|
||||||
|
const poolScoreHardStateOptions = [
|
||||||
|
{ value: 'all', label: '全部状态' },
|
||||||
|
{ value: 'available', label: '可用' },
|
||||||
|
{ value: 'unknown', label: '未知' },
|
||||||
|
{ value: 'cooldown', label: '冷却' },
|
||||||
|
{ value: 'quota_exhausted', label: '额度耗尽' },
|
||||||
|
{ value: 'auth_invalid', label: '授权无效' },
|
||||||
|
{ value: 'banned', label: '封禁' },
|
||||||
|
{ value: 'inactive', label: '禁用' },
|
||||||
|
]
|
||||||
|
const poolScoreProbeStatusOptions = [
|
||||||
|
{ value: 'all', label: '全部探测' },
|
||||||
|
{ value: 'never', label: '未探测' },
|
||||||
|
{ value: 'ok', label: '正常' },
|
||||||
|
{ value: 'failed', label: '失败' },
|
||||||
|
{ value: 'stale', label: '过期' },
|
||||||
|
{ value: 'in_progress', label: '探测中' },
|
||||||
|
]
|
||||||
|
|
||||||
async function loadOverview(options: { cacheTtlMs?: number } = {}) {
|
async function loadOverview(options: { cacheTtlMs?: number } = {}) {
|
||||||
const requestId = ++overviewRequestId
|
const requestId = ++overviewRequestId
|
||||||
@@ -1488,6 +1676,7 @@ async function loadOverview(options: { cacheTtlMs?: number } = {}) {
|
|||||||
showAccountBatchDialog.value = false
|
showAccountBatchDialog.value = false
|
||||||
closeProviderProxyPopovers()
|
closeProviderProxyPopovers()
|
||||||
resetKeyPage()
|
resetKeyPage()
|
||||||
|
resetScorePage()
|
||||||
}
|
}
|
||||||
} catch (err) {
|
} catch (err) {
|
||||||
if (requestId !== overviewRequestId) return
|
if (requestId !== overviewRequestId) return
|
||||||
@@ -1720,9 +1909,11 @@ async function selectProvider(
|
|||||||
}
|
}
|
||||||
keysLoadedOnce.value = false
|
keysLoadedOnce.value = false
|
||||||
resetKeyPage(currentPage.value, pageSize.value)
|
resetKeyPage(currentPage.value, pageSize.value)
|
||||||
|
resetScorePage()
|
||||||
const keysTask = loadKeys({ cacheTtlMs: options.cacheTtlMs ?? 0 })
|
const keysTask = loadKeys({ cacheTtlMs: options.cacheTtlMs ?? 0 })
|
||||||
// Provider summary is non-blocking for key list rendering.
|
// Provider summary is non-blocking for key list rendering.
|
||||||
void loadProviderData(id)
|
void loadProviderData(id)
|
||||||
|
void loadScores({ cacheTtlMs: options.cacheTtlMs ? POOL_SCORES_CACHE_TTL_MS : 0 })
|
||||||
await keysTask
|
await keysTask
|
||||||
if (requestId !== selectProviderRequestId) return
|
if (requestId !== selectProviderRequestId) return
|
||||||
}
|
}
|
||||||
@@ -1748,18 +1939,40 @@ function createEmptyKeyPage(page = 1, pageSizeValue = 50): PoolKeysPageResponse
|
|||||||
return { total: 0, page, page_size: pageSizeValue, keys: [] }
|
return { total: 0, page, page_size: pageSizeValue, keys: [] }
|
||||||
}
|
}
|
||||||
|
|
||||||
|
function createEmptyScorePage(): PoolScoresResponse {
|
||||||
|
return {
|
||||||
|
provider_id: selectedProviderId.value || '',
|
||||||
|
page: 1,
|
||||||
|
page_size: 20,
|
||||||
|
filters: {},
|
||||||
|
items: [],
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
const keyPage = ref<PoolKeysPageResponse>(createEmptyKeyPage())
|
const keyPage = ref<PoolKeysPageResponse>(createEmptyKeyPage())
|
||||||
|
const poolScorePage = ref<PoolScoresResponse>(createEmptyScorePage())
|
||||||
const keysLoading = ref(false)
|
const keysLoading = ref(false)
|
||||||
|
const scoresLoading = ref(false)
|
||||||
const keysLoadedOnce = ref(false)
|
const keysLoadedOnce = ref(false)
|
||||||
const refreshingCurrentPageQuota = ref(false)
|
const refreshingCurrentPageQuota = ref(false)
|
||||||
const searchQuery = ref(restoredViewState.search)
|
const searchQuery = ref(restoredViewState.search)
|
||||||
const statusFilter = ref(restoredViewState.status)
|
const statusFilter = ref(restoredViewState.status)
|
||||||
|
const scoreApiFormatFilter = ref('')
|
||||||
|
const scoreModelIdFilter = ref('')
|
||||||
|
const scoreHardStateFilter = ref('all')
|
||||||
|
const scoreProbeStatusFilter = ref('all')
|
||||||
const currentPage = ref(restoredViewState.page)
|
const currentPage = ref(restoredViewState.page)
|
||||||
const pageSize = ref(restoredViewState.pageSize)
|
const pageSize = ref(restoredViewState.pageSize)
|
||||||
const sortBy = ref<PoolManagementSortBy | null>(restoredViewState.sortBy)
|
const sortBy = ref<PoolManagementSortBy | null>(restoredViewState.sortBy)
|
||||||
const sortOrder = ref<PoolManagementSortOrder>(restoredViewState.sortOrder)
|
const sortOrder = ref<PoolManagementSortOrder>(restoredViewState.sortOrder)
|
||||||
const poolStatsMode = ref<PoolManagementStatsMode>(restoredViewState.statsMode)
|
const poolStatsMode = ref<PoolManagementStatsMode>(restoredViewState.statsMode)
|
||||||
const hasPoolKeyFilters = computed(() => searchQuery.value.trim().length > 0 || statusFilter.value !== 'all')
|
const hasPoolKeyFilters = computed(() => searchQuery.value.trim().length > 0 || statusFilter.value !== 'all')
|
||||||
|
const hasPoolScoreFilters = computed(() => {
|
||||||
|
return scoreApiFormatFilter.value.trim().length > 0
|
||||||
|
|| scoreModelIdFilter.value.trim().length > 0
|
||||||
|
|| scoreHardStateFilter.value !== 'all'
|
||||||
|
|| scoreProbeStatusFilter.value !== 'all'
|
||||||
|
})
|
||||||
const MANUAL_QUOTA_REFRESH_COOLDOWN_SECONDS = 5 * 60
|
const MANUAL_QUOTA_REFRESH_COOLDOWN_SECONDS = 5 * 60
|
||||||
const refreshingOAuthKeyId = ref<string | null>(null)
|
const refreshingOAuthKeyId = ref<string | null>(null)
|
||||||
const resettingCycleKeyId = ref<string | null>(null)
|
const resettingCycleKeyId = ref<string | null>(null)
|
||||||
@@ -2060,6 +2273,10 @@ function resetKeyPage(page = currentPage.value, pageSizeValue = pageSize.value):
|
|||||||
keyPage.value = createEmptyKeyPage(page, pageSizeValue)
|
keyPage.value = createEmptyKeyPage(page, pageSizeValue)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
function resetScorePage(): void {
|
||||||
|
poolScorePage.value = createEmptyScorePage()
|
||||||
|
}
|
||||||
|
|
||||||
function refreshOverviewInBackground(): void {
|
function refreshOverviewInBackground(): void {
|
||||||
void loadOverview()
|
void loadOverview()
|
||||||
}
|
}
|
||||||
@@ -2199,6 +2416,36 @@ async function refreshCurrentPage() {
|
|||||||
if (!quotaDidReload) {
|
if (!quotaDidReload) {
|
||||||
await refresh()
|
await refresh()
|
||||||
}
|
}
|
||||||
|
void loadScores()
|
||||||
|
}
|
||||||
|
|
||||||
|
async function loadScores(options: { cacheTtlMs?: number } = {}) {
|
||||||
|
if (!selectedProviderId.value) return
|
||||||
|
const requestId = ++scoresRequestId
|
||||||
|
const providerId = selectedProviderId.value
|
||||||
|
scoresLoading.value = true
|
||||||
|
try {
|
||||||
|
const nextPage = await listPoolScores(providerId, {
|
||||||
|
page: 1,
|
||||||
|
page_size: 20,
|
||||||
|
api_format: scoreApiFormatFilter.value.trim() || undefined,
|
||||||
|
model_id: scoreModelIdFilter.value.trim() || undefined,
|
||||||
|
hard_state: scoreHardStateFilter.value === 'all' ? undefined : scoreHardStateFilter.value,
|
||||||
|
probe_status: scoreProbeStatusFilter.value === 'all' ? undefined : scoreProbeStatusFilter.value,
|
||||||
|
}, {
|
||||||
|
cacheTtlMs: options.cacheTtlMs ?? 0,
|
||||||
|
})
|
||||||
|
if (requestId !== scoresRequestId || selectedProviderId.value !== providerId) return
|
||||||
|
poolScorePage.value = nextPage
|
||||||
|
} catch (err) {
|
||||||
|
if (requestId !== scoresRequestId || selectedProviderId.value !== providerId) return
|
||||||
|
resetScorePage()
|
||||||
|
showError(parseApiError(err, '读取评分失败'))
|
||||||
|
} finally {
|
||||||
|
if (requestId === scoresRequestId) {
|
||||||
|
scoresLoading.value = false
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
async function loadKeys(options: { cacheTtlMs?: number } = {}) {
|
async function loadKeys(options: { cacheTtlMs?: number } = {}) {
|
||||||
@@ -2276,6 +2523,20 @@ watch(searchQuery, () => {
|
|||||||
}, 300)
|
}, 300)
|
||||||
})
|
})
|
||||||
|
|
||||||
|
watch([scoreApiFormatFilter, scoreModelIdFilter], () => {
|
||||||
|
if (scoresFilterDebounceTimer !== null) {
|
||||||
|
clearTimeout(scoresFilterDebounceTimer)
|
||||||
|
}
|
||||||
|
scoresFilterDebounceTimer = window.setTimeout(() => {
|
||||||
|
scoresFilterDebounceTimer = null
|
||||||
|
void loadScores({ cacheTtlMs: POOL_SCORES_CACHE_TTL_MS })
|
||||||
|
}, 300)
|
||||||
|
})
|
||||||
|
|
||||||
|
watch([scoreHardStateFilter, scoreProbeStatusFilter], () => {
|
||||||
|
void loadScores({ cacheTtlMs: POOL_SCORES_CACHE_TTL_MS })
|
||||||
|
})
|
||||||
|
|
||||||
function normalizeAuthTypeForEdit(key: PoolKeyDetail): EndpointAPIKey['auth_type'] {
|
function normalizeAuthTypeForEdit(key: PoolKeyDetail): EndpointAPIKey['auth_type'] {
|
||||||
if (isOAuthManagedCredential(key)) return 'oauth'
|
if (isOAuthManagedCredential(key)) return 'oauth'
|
||||||
if (isServiceAccountCredential(key)) return 'service_account'
|
if (isServiceAccountCredential(key)) return 'service_account'
|
||||||
@@ -3584,6 +3845,26 @@ function formatStatUsd(value: number | string | null | undefined): string {
|
|||||||
return `$${n.toLocaleString('en-US', { minimumFractionDigits: 2, maximumFractionDigits: 2 })}`
|
return `$${n.toLocaleString('en-US', { minimumFractionDigits: 2, maximumFractionDigits: 2 })}`
|
||||||
}
|
}
|
||||||
|
|
||||||
|
function formatPoolScore(value: number): string {
|
||||||
|
const n = Number(value)
|
||||||
|
if (!Number.isFinite(n)) return '0.000'
|
||||||
|
return n.toFixed(3)
|
||||||
|
}
|
||||||
|
|
||||||
|
function getPoolScoreHardStateLabel(value: PoolMemberScoreItem['hard_state']): string {
|
||||||
|
return poolScoreHardStateOptions.find(item => item.value === value)?.label || value
|
||||||
|
}
|
||||||
|
|
||||||
|
function getPoolScoreProbeStatusLabel(value: PoolMemberScoreItem['probe_status']): string {
|
||||||
|
return poolScoreProbeStatusOptions.find(item => item.value === value)?.label || value
|
||||||
|
}
|
||||||
|
|
||||||
|
function formatUnixSeconds(seconds: number | null | undefined): string {
|
||||||
|
const raw = Number(seconds ?? 0)
|
||||||
|
if (!Number.isFinite(raw) || raw <= 0) return '-'
|
||||||
|
return formatRelativeTime(new Date(raw * 1000).toISOString())
|
||||||
|
}
|
||||||
|
|
||||||
function formatRelativeTime(isoStr: string): string {
|
function formatRelativeTime(isoStr: string): string {
|
||||||
const date = new Date(isoStr)
|
const date = new Date(isoStr)
|
||||||
const pad = (n: number) => String(n).padStart(2, '0')
|
const pad = (n: number) => String(n).padStart(2, '0')
|
||||||
@@ -3611,9 +3892,14 @@ onBeforeUnmount(() => {
|
|||||||
clearTimeout(keysSearchDebounceTimer)
|
clearTimeout(keysSearchDebounceTimer)
|
||||||
keysSearchDebounceTimer = null
|
keysSearchDebounceTimer = null
|
||||||
}
|
}
|
||||||
|
if (scoresFilterDebounceTimer !== null) {
|
||||||
|
clearTimeout(scoresFilterDebounceTimer)
|
||||||
|
scoresFilterDebounceTimer = null
|
||||||
|
}
|
||||||
overviewRequestId += 1
|
overviewRequestId += 1
|
||||||
selectProviderRequestId += 1
|
selectProviderRequestId += 1
|
||||||
providerDataRequestId += 1
|
providerDataRequestId += 1
|
||||||
keysRequestId += 1
|
keysRequestId += 1
|
||||||
|
scoresRequestId += 1
|
||||||
})
|
})
|
||||||
</script>
|
</script>
|
||||||
|
|||||||
Reference in New Issue
Block a user