mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
Merge remote-tracking branch 'origin/aether-rust-pioneer' into payment-billing-plans
# Conflicts: # crates/aether-data/src/lifecycle/bootstrap/postgres.rs # crates/aether-data/src/lifecycle/migrate/tests.rs
This commit is contained in:
@@ -9,6 +9,7 @@ description = "AI serving application contracts and ports for Aether"
|
||||
[dependencies]
|
||||
aether-ai-formats.workspace = true
|
||||
aether-contracts.workspace = true
|
||||
aether-data-contracts.workspace = true
|
||||
aether-scheduler-core.workspace = true
|
||||
async-trait.workspace = true
|
||||
http.workspace = true
|
||||
|
||||
@@ -20,7 +20,7 @@ impl<'a> AiCandidateResolutionRequest<'a> {
|
||||
client_api_format,
|
||||
requested_model,
|
||||
mode: AiCandidateResolutionMode::Standard,
|
||||
expand_pool_groups: true,
|
||||
expand_pool_groups: false,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -32,7 +32,7 @@ impl<'a> AiCandidateResolutionRequest<'a> {
|
||||
client_api_format,
|
||||
requested_model,
|
||||
mode: AiCandidateResolutionMode::WithoutTransportPairGate,
|
||||
expand_pool_groups: true,
|
||||
expand_pool_groups: false,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -328,7 +328,7 @@ mod tests {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn resolution_reads_transport_gates_candidates_then_ranks_and_applies_pool() {
|
||||
async fn resolution_reads_transport_gates_candidates_then_ranks_without_expanding_pools() {
|
||||
let port = TestPort::default();
|
||||
|
||||
let outcome = run_ai_candidate_resolution(
|
||||
@@ -349,7 +349,6 @@ mod tests {
|
||||
"missing:transport_snapshot_missing",
|
||||
"inactive:provider_inactive",
|
||||
"unsupported:transport_unsupported",
|
||||
"pool:cooldown",
|
||||
]
|
||||
);
|
||||
assert_eq!(
|
||||
@@ -368,7 +367,6 @@ mod tests {
|
||||
"common:second:gpt-4.1",
|
||||
"pair:second:openai:chat:gpt-4.1",
|
||||
"rank:openai:chat",
|
||||
"pool",
|
||||
]
|
||||
);
|
||||
}
|
||||
@@ -386,14 +384,13 @@ mod tests {
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(outcome.eligible_candidates, ["eligible:unsupported"]);
|
||||
assert_eq!(outcome.skipped_candidates, ["pool:cooldown"]);
|
||||
assert!(outcome.skipped_candidates.is_empty());
|
||||
assert_eq!(
|
||||
port.calls.lock().unwrap().as_slice(),
|
||||
[
|
||||
"transport:unsupported",
|
||||
"common:unsupported:",
|
||||
"rank:openai:chat",
|
||||
"pool",
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
@@ -16,6 +16,7 @@ pub mod execution_path;
|
||||
pub mod failure_diagnostic;
|
||||
pub mod plan_payload;
|
||||
pub mod pool_scheduler;
|
||||
pub mod pool_scores;
|
||||
pub mod ports;
|
||||
pub mod ranking_metadata;
|
||||
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_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 report_context::{
|
||||
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::{
|
||||
MinimalCandidateSelectionReadRepository, MinimalCandidateSelectionRepository,
|
||||
StoredMinimalCandidateSelectionRow, StoredPoolKeyCandidateOrder,
|
||||
StoredPoolKeyCandidateRowsQuery, StoredProviderModelMapping,
|
||||
StoredRequestedModelCandidateRowsQuery,
|
||||
StoredPoolKeyCandidateRowsByKeyIdsQuery, StoredPoolKeyCandidateRowsQuery,
|
||||
StoredProviderModelMapping, StoredRequestedModelCandidateRowsQuery,
|
||||
};
|
||||
|
||||
@@ -67,6 +67,16 @@ pub struct StoredPoolKeyCandidateRowsQuery {
|
||||
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)]
|
||||
pub struct StoredRequestedModelCandidateRowsQuery {
|
||||
pub api_format: String,
|
||||
@@ -124,6 +134,11 @@ pub trait MinimalCandidateSelectionReadRepository: Send + Sync {
|
||||
&self,
|
||||
query: &StoredPoolKeyCandidateRowsQuery,
|
||||
) -> 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:
|
||||
|
||||
@@ -3,6 +3,7 @@ pub mod billing;
|
||||
pub mod candidate_selection;
|
||||
pub mod candidates;
|
||||
pub mod global_models;
|
||||
pub mod pool_scores;
|
||||
pub mod provider_catalog;
|
||||
pub mod quota;
|
||||
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_ACCOUNT,
|
||||
POOL_SCORE_CAPABILITY_API_FORMAT, POOL_SCORE_SCOPE_KIND_ACCOUNT, POOL_SCORE_SCOPE_KIND_MODEL,
|
||||
};
|
||||
361
crates/aether-data-contracts/src/repository/pool_scores/types.rs
Normal file
361
crates/aether-data-contracts/src/repository/pool_scores/types.rs
Normal file
@@ -0,0 +1,361 @@
|
||||
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_ACCOUNT: &str = "account";
|
||||
pub const POOL_SCORE_SCOPE_KIND_ACCOUNT: &str = "account";
|
||||
pub const POOL_SCORE_CAPABILITY_API_FORMAT: &str = POOL_SCORE_CAPABILITY_ACCOUNT;
|
||||
pub const POOL_SCORE_SCOPE_KIND_MODEL: &str = POOL_SCORE_SCOPE_KIND_ACCOUNT;
|
||||
|
||||
#[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)
|
||||
);
|
||||
|
||||
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 (
|
||||
id VARCHAR(64) PRIMARY KEY,
|
||||
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: -
|
||||
--
|
||||
@@ -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: -
|
||||
--
|
||||
@@ -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: -
|
||||
--
|
||||
|
||||
@@ -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 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 (
|
||||
id TEXT PRIMARY KEY,
|
||||
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);
|
||||
@@ -625,6 +625,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: -
|
||||
--
|
||||
|
||||
@@ -297,6 +297,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: -
|
||||
--
|
||||
|
||||
@@ -925,6 +925,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: -
|
||||
--
|
||||
|
||||
@@ -127,6 +127,36 @@ CREATE TABLE IF NOT EXISTS provider_api_keys (
|
||||
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 (
|
||||
id VARCHAR(64) PRIMARY KEY,
|
||||
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: -
|
||||
--
|
||||
|
||||
@@ -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: -
|
||||
--
|
||||
|
||||
@@ -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: -
|
||||
--
|
||||
|
||||
@@ -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 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 (
|
||||
id TEXT PRIMARY KEY,
|
||||
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`)
|
||||
);
|
||||
|
||||
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 (
|
||||
`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);
|
||||
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 (
|
||||
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 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 (
|
||||
id TEXT PRIMARY KEY NOT NULL,
|
||||
api_key_id TEXT NOT NULL,
|
||||
|
||||
@@ -579,6 +579,143 @@ type = "unix_seconds"
|
||||
name = "provider_api_keys_provider_id_idx"
|
||||
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]
|
||||
domain = "provider_catalog"
|
||||
order = 45
|
||||
|
||||
@@ -37,6 +37,9 @@ use crate::repository::management_tokens::{
|
||||
use crate::repository::oauth_providers::{
|
||||
MysqlOAuthProviderRepository, OAuthProviderReadRepository, OAuthProviderWriteRepository,
|
||||
};
|
||||
use crate::repository::pool_scores::{
|
||||
MysqlPoolMemberScoreRepository, PoolMemberScoreWriteRepository, PoolScoreReadRepository,
|
||||
};
|
||||
use crate::repository::provider_catalog::{
|
||||
MysqlProviderCatalogReadRepository, ProviderCatalogReadRepository,
|
||||
ProviderCatalogWriteRepository,
|
||||
@@ -184,6 +187,14 @@ impl MysqlBackend {
|
||||
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> {
|
||||
Arc::new(MysqlProxyNodeReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
@@ -39,6 +39,9 @@ use crate::repository::management_tokens::{
|
||||
use crate::repository::oauth_providers::{
|
||||
OAuthProviderReadRepository, OAuthProviderWriteRepository, SqlxOAuthProviderRepository,
|
||||
};
|
||||
use crate::repository::pool_scores::{
|
||||
PoolMemberScoreWriteRepository, PoolScoreReadRepository, PostgresPoolMemberScoreRepository,
|
||||
};
|
||||
use crate::repository::provider_catalog::{
|
||||
ProviderCatalogReadRepository, ProviderCatalogWriteRepository,
|
||||
SqlxProviderCatalogReadRepository,
|
||||
@@ -195,6 +198,14 @@ impl PostgresBackend {
|
||||
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> {
|
||||
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::management_tokens::ManagementTokenReadRepository;
|
||||
use crate::repository::oauth_providers::OAuthProviderReadRepository;
|
||||
use crate::repository::pool_scores::PoolScoreReadRepository;
|
||||
use crate::repository::provider_catalog::ProviderCatalogReadRepository;
|
||||
use crate::repository::proxy_nodes::ProxyNodeReadRepository;
|
||||
use crate::repository::quota::ProviderQuotaReadRepository;
|
||||
@@ -34,6 +35,7 @@ pub struct DataReadRepositories {
|
||||
global_models: Option<Arc<dyn GlobalModelReadRepository>>,
|
||||
management_tokens: Option<Arc<dyn ManagementTokenReadRepository>>,
|
||||
oauth_providers: Option<Arc<dyn OAuthProviderReadRepository>>,
|
||||
pool_scores: Option<Arc<dyn PoolScoreReadRepository>>,
|
||||
proxy_nodes: Option<Arc<dyn ProxyNodeReadRepository>>,
|
||||
minimal_candidate_selection: Option<Arc<dyn MinimalCandidateSelectionReadRepository>>,
|
||||
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_management_tokens", &self.management_tokens.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_minimal_candidate_selection",
|
||||
@@ -124,6 +127,10 @@ impl DataReadRepositories {
|
||||
.map(PostgresBackend::oauth_provider_read_repository)
|
||||
.or_else(|| mysql.map(MysqlBackend::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
|
||||
.map(PostgresBackend::proxy_node_read_repository)
|
||||
.or_else(|| mysql.map(MysqlBackend::proxy_node_read_repository))
|
||||
@@ -208,6 +215,10 @@ impl DataReadRepositories {
|
||||
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>> {
|
||||
self.proxy_nodes.clone()
|
||||
}
|
||||
@@ -257,6 +268,7 @@ impl DataReadRepositories {
|
||||
|| self.global_models.is_some()
|
||||
|| self.management_tokens.is_some()
|
||||
|| self.oauth_providers.is_some()
|
||||
|| self.pool_scores.is_some()
|
||||
|| self.proxy_nodes.is_some()
|
||||
|| self.minimal_candidate_selection.is_some()
|
||||
|| self.request_candidates.is_some()
|
||||
|
||||
@@ -37,6 +37,9 @@ use crate::repository::management_tokens::{
|
||||
use crate::repository::oauth_providers::{
|
||||
OAuthProviderReadRepository, OAuthProviderWriteRepository, SqliteOAuthProviderRepository,
|
||||
};
|
||||
use crate::repository::pool_scores::{
|
||||
PoolMemberScoreWriteRepository, PoolScoreReadRepository, SqlitePoolMemberScoreRepository,
|
||||
};
|
||||
use crate::repository::provider_catalog::{
|
||||
ProviderCatalogReadRepository, ProviderCatalogWriteRepository,
|
||||
SqliteProviderCatalogReadRepository,
|
||||
@@ -197,6 +200,14 @@ impl SqliteBackend {
|
||||
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> {
|
||||
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::management_tokens::ManagementTokenWriteRepository;
|
||||
use crate::repository::oauth_providers::OAuthProviderWriteRepository;
|
||||
use crate::repository::pool_scores::PoolMemberScoreWriteRepository;
|
||||
use crate::repository::provider_catalog::ProviderCatalogWriteRepository;
|
||||
use crate::repository::proxy_nodes::ProxyNodeWriteRepository;
|
||||
use crate::repository::quota::ProviderQuotaWriteRepository;
|
||||
@@ -30,6 +31,7 @@ pub struct DataWriteRepositories {
|
||||
global_models: Option<Arc<dyn GlobalModelWriteRepository>>,
|
||||
management_tokens: Option<Arc<dyn ManagementTokenWriteRepository>>,
|
||||
oauth_providers: Option<Arc<dyn OAuthProviderWriteRepository>>,
|
||||
pool_scores: Option<Arc<dyn PoolMemberScoreWriteRepository>>,
|
||||
proxy_nodes: Option<Arc<dyn ProxyNodeWriteRepository>>,
|
||||
provider_catalog: Option<Arc<dyn ProviderCatalogWriteRepository>>,
|
||||
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_management_tokens", &self.management_tokens.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_provider_catalog", &self.provider_catalog.is_some())
|
||||
.field("has_provider_quotas", &self.provider_quotas.is_some())
|
||||
@@ -108,6 +111,10 @@ impl DataWriteRepositories {
|
||||
.map(PostgresBackend::oauth_provider_write_repository)
|
||||
.or_else(|| mysql.map(MysqlBackend::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
|
||||
.map(PostgresBackend::proxy_node_write_repository)
|
||||
.or_else(|| mysql.map(MysqlBackend::proxy_node_write_repository))
|
||||
@@ -184,6 +191,10 @@ impl DataWriteRepositories {
|
||||
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>> {
|
||||
self.proxy_nodes.clone()
|
||||
}
|
||||
@@ -218,6 +229,7 @@ impl DataWriteRepositories {
|
||||
|| self.global_models.is_some()
|
||||
|| self.management_tokens.is_some()
|
||||
|| self.oauth_providers.is_some()
|
||||
|| self.pool_scores.is_some()
|
||||
|| self.proxy_nodes.is_some()
|
||||
|| self.provider_catalog.is_some()
|
||||
|| self.provider_quotas.is_some()
|
||||
|
||||
@@ -300,6 +300,7 @@ fn empty_database_snapshot_covers_current_cutoff_versions() {
|
||||
20260510120000,
|
||||
20260511000000,
|
||||
20260511120000,
|
||||
20260511130000,
|
||||
20260512090000,
|
||||
20260512110000,
|
||||
]
|
||||
@@ -579,8 +580,9 @@ fn mysql_and_sqlite_migrations_include_enabled_incrementals() {
|
||||
20260509120000,
|
||||
20260510120000,
|
||||
20260511120000,
|
||||
20260511130000,
|
||||
20260512090000,
|
||||
20260512110000
|
||||
20260512110000,
|
||||
]
|
||||
);
|
||||
assert_eq!(
|
||||
@@ -593,8 +595,9 @@ fn mysql_and_sqlite_migrations_include_enabled_incrementals() {
|
||||
20260509120000,
|
||||
20260510120000,
|
||||
20260511120000,
|
||||
20260511130000,
|
||||
20260512090000,
|
||||
20260512110000
|
||||
20260512110000,
|
||||
]
|
||||
);
|
||||
}
|
||||
@@ -1106,6 +1109,7 @@ fn pending_migrations_from_applied_skips_versions_already_applied() {
|
||||
20260510120000,
|
||||
20260511000000,
|
||||
20260511120000,
|
||||
20260511130000,
|
||||
20260512090000,
|
||||
20260512110000,
|
||||
]
|
||||
|
||||
@@ -4,8 +4,8 @@ use async_trait::async_trait;
|
||||
|
||||
use super::{
|
||||
MinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
|
||||
StoredPoolKeyCandidateOrder, StoredPoolKeyCandidateRowsQuery,
|
||||
StoredRequestedModelCandidateRowsQuery,
|
||||
StoredPoolKeyCandidateOrder, StoredPoolKeyCandidateRowsByKeyIdsQuery,
|
||||
StoredPoolKeyCandidateRowsQuery, StoredRequestedModelCandidateRowsQuery,
|
||||
};
|
||||
use crate::DataLayerError;
|
||||
|
||||
@@ -138,6 +138,39 @@ impl MinimalCandidateSelectionReadRepository for InMemoryMinimalCandidateSelecti
|
||||
.take(query.limit as usize)
|
||||
.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(
|
||||
@@ -185,8 +218,10 @@ fn row_matches_requested_model(
|
||||
requested_model_name: &str,
|
||||
api_format: &str,
|
||||
) -> bool {
|
||||
row.global_model_name == requested_model_name
|
||||
|| row.model_provider_model_name == requested_model_name
|
||||
(row_has_available_provider_model(row, api_format)
|
||||
&& row.global_model_name == requested_model_name)
|
||||
|| (row_default_provider_model_name_available(row, api_format)
|
||||
&& row.model_provider_model_name == requested_model_name)
|
||||
|| row
|
||||
.model_provider_model_mappings
|
||||
.as_ref()
|
||||
@@ -205,6 +240,60 @@ fn row_matches_requested_model(
|
||||
})
|
||||
}
|
||||
|
||||
fn row_has_available_provider_model(
|
||||
row: &StoredMinimalCandidateSelectionRow,
|
||||
api_format: &str,
|
||||
) -> bool {
|
||||
row_mapping_matches_scope(row, api_format)
|
||||
|| row_default_provider_model_name_available(row, api_format)
|
||||
}
|
||||
|
||||
fn row_default_provider_model_name_available(
|
||||
row: &StoredMinimalCandidateSelectionRow,
|
||||
api_format: &str,
|
||||
) -> bool {
|
||||
let Some(mappings) = row.model_provider_model_mappings.as_ref() else {
|
||||
return true;
|
||||
};
|
||||
let mut has_explicit_default_mapping = false;
|
||||
for mapping in mappings {
|
||||
if mapping.name != row.model_provider_model_name {
|
||||
continue;
|
||||
}
|
||||
has_explicit_default_mapping = true;
|
||||
if mapping_scope_matches(mapping, row, api_format) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
!has_explicit_default_mapping
|
||||
}
|
||||
|
||||
fn row_mapping_matches_scope(row: &StoredMinimalCandidateSelectionRow, api_format: &str) -> bool {
|
||||
row.model_provider_model_mappings
|
||||
.as_ref()
|
||||
.is_some_and(|mappings| {
|
||||
mappings
|
||||
.iter()
|
||||
.any(|mapping| mapping_scope_matches(mapping, row, api_format))
|
||||
})
|
||||
}
|
||||
|
||||
fn mapping_scope_matches(
|
||||
mapping: &super::StoredProviderModelMapping,
|
||||
row: &StoredMinimalCandidateSelectionRow,
|
||||
api_format: &str,
|
||||
) -> bool {
|
||||
mapping.api_formats.as_ref().is_none_or(|formats| {
|
||||
formats
|
||||
.iter()
|
||||
.any(|value| api_format_matches(value, api_format))
|
||||
}) && mapping.endpoint_ids.as_ref().is_none_or(|endpoint_ids| {
|
||||
endpoint_ids
|
||||
.iter()
|
||||
.any(|endpoint_id| endpoint_id == &row.endpoint_id)
|
||||
})
|
||||
}
|
||||
|
||||
fn key_auth_channel_matches(row: &StoredMinimalCandidateSelectionRow, api_format: &str) -> bool {
|
||||
let provider_type = row.provider_type.trim().to_ascii_lowercase();
|
||||
let auth_type = row.key_auth_type.trim().to_ascii_lowercase();
|
||||
@@ -244,7 +333,7 @@ mod tests {
|
||||
use super::InMemoryMinimalCandidateSelectionReadRepository;
|
||||
use crate::repository::candidate_selection::{
|
||||
MinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
|
||||
StoredPoolKeyCandidateOrder, StoredPoolKeyCandidateRowsQuery,
|
||||
StoredPoolKeyCandidateOrder, StoredPoolKeyCandidateRowsQuery, StoredProviderModelMapping,
|
||||
StoredRequestedModelCandidateRowsQuery,
|
||||
};
|
||||
|
||||
@@ -310,14 +399,12 @@ mod tests {
|
||||
async fn filters_by_exact_api_format_and_requested_model_aliases() {
|
||||
let mut mapped = sample_row("provider-1", "openai:chat", "gpt-4.1", 10);
|
||||
mapped.model_provider_model_name = "provider-gpt-4.1".to_string();
|
||||
mapped.model_provider_model_mappings = Some(vec![
|
||||
crate::repository::candidate_selection::StoredProviderModelMapping {
|
||||
name: "alias-gpt-4.1".to_string(),
|
||||
priority: 0,
|
||||
api_formats: Some(vec!["openai:chat".to_string()]),
|
||||
endpoint_ids: None,
|
||||
},
|
||||
]);
|
||||
mapped.model_provider_model_mappings = Some(vec![StoredProviderModelMapping {
|
||||
name: "alias-gpt-4.1".to_string(),
|
||||
priority: 0,
|
||||
api_formats: Some(vec!["openai:chat".to_string()]),
|
||||
endpoint_ids: None,
|
||||
}]);
|
||||
let repository = InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
|
||||
mapped,
|
||||
sample_row("provider-2", "openai:chat", "gpt-4.1-mini", 20),
|
||||
@@ -332,6 +419,42 @@ mod tests {
|
||||
assert_eq!(rows[0].provider_id, "provider-1");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn requested_model_filter_respects_endpoint_scoped_default_mapping() {
|
||||
let mut selected = sample_row("provider-1", "openai:chat", "deepseek-v4-pro", 10);
|
||||
selected.endpoint_id = "endpoint-openai".to_string();
|
||||
selected.model_provider_model_name = "deepseek-v4-pro".to_string();
|
||||
selected.model_provider_model_mappings = Some(vec![StoredProviderModelMapping {
|
||||
name: "deepseek-v4-pro".to_string(),
|
||||
priority: 1,
|
||||
api_formats: None,
|
||||
endpoint_ids: Some(vec!["endpoint-openai".to_string()]),
|
||||
}]);
|
||||
|
||||
let mut scoped_out = selected.clone();
|
||||
scoped_out.provider_id = "provider-2".to_string();
|
||||
scoped_out.endpoint_id = "endpoint-claude".to_string();
|
||||
scoped_out.endpoint_api_format = "claude:messages".to_string();
|
||||
scoped_out.key_id = "key-provider-2".to_string();
|
||||
scoped_out.key_api_formats = Some(vec!["claude:messages".to_string()]);
|
||||
|
||||
let repository =
|
||||
InMemoryMinimalCandidateSelectionReadRepository::seed(vec![scoped_out, selected]);
|
||||
|
||||
let rows = repository
|
||||
.list_for_exact_api_format_and_requested_model("claude:messages", "deepseek-v4-pro")
|
||||
.await
|
||||
.expect("list should succeed");
|
||||
assert!(rows.is_empty());
|
||||
|
||||
let rows = repository
|
||||
.list_for_exact_api_format_and_requested_model("openai:chat", "deepseek-v4-pro")
|
||||
.await
|
||||
.expect("list should succeed");
|
||||
assert_eq!(rows.len(), 1);
|
||||
assert_eq!(rows[0].endpoint_id, "endpoint-openai");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn requested_model_page_returns_requested_slice_only() {
|
||||
let mut rows = Vec::new();
|
||||
|
||||
@@ -7,8 +7,8 @@ mod sqlite;
|
||||
pub(crate) use aether_data_contracts::repository::candidate_selection::{
|
||||
MinimalCandidateSelectionReadRepository, MinimalCandidateSelectionRepository,
|
||||
StoredMinimalCandidateSelectionRow, StoredPoolKeyCandidateOrder,
|
||||
StoredPoolKeyCandidateRowsQuery, StoredProviderModelMapping,
|
||||
StoredRequestedModelCandidateRowsQuery,
|
||||
StoredPoolKeyCandidateRowsByKeyIdsQuery, StoredPoolKeyCandidateRowsQuery,
|
||||
StoredProviderModelMapping, StoredRequestedModelCandidateRowsQuery,
|
||||
};
|
||||
pub use memory::InMemoryMinimalCandidateSelectionReadRepository;
|
||||
pub use mysql::MysqlMinimalCandidateSelectionReadRepository;
|
||||
|
||||
@@ -5,7 +5,8 @@ use sqlx::{mysql::MySqlRow, MySql, QueryBuilder, Row};
|
||||
|
||||
use super::{
|
||||
MinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
|
||||
StoredPoolKeyCandidateOrder, StoredPoolKeyCandidateRowsQuery, StoredProviderModelMapping,
|
||||
StoredPoolKeyCandidateOrder, StoredPoolKeyCandidateRowsByKeyIdsQuery,
|
||||
StoredPoolKeyCandidateRowsQuery, StoredProviderModelMapping,
|
||||
StoredRequestedModelCandidateRowsQuery,
|
||||
};
|
||||
use crate::driver::mysql::MysqlPool;
|
||||
@@ -196,6 +197,40 @@ impl MinimalCandidateSelectionReadRepository for MysqlMinimalCandidateSelectionR
|
||||
.map(|item| item.row)
|
||||
.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> {
|
||||
@@ -313,26 +348,75 @@ fn row_matches_requested_model(
|
||||
requested_model_name: &str,
|
||||
api_format: &str,
|
||||
) -> bool {
|
||||
row.global_model_name == requested_model_name
|
||||
|| row.model_provider_model_name == requested_model_name
|
||||
(row_has_available_provider_model(row, api_format)
|
||||
&& row.global_model_name == requested_model_name)
|
||||
|| (row_default_provider_model_name_available(row, api_format)
|
||||
&& row.model_provider_model_name == requested_model_name)
|
||||
|| row
|
||||
.model_provider_model_mappings
|
||||
.as_ref()
|
||||
.is_some_and(|mappings| {
|
||||
mappings.iter().any(|mapping| {
|
||||
mapping.api_formats.as_ref().is_none_or(|formats| {
|
||||
formats
|
||||
.iter()
|
||||
.any(|value| api_format_matches(value, api_format))
|
||||
}) && mapping.endpoint_ids.as_ref().is_none_or(|endpoint_ids| {
|
||||
endpoint_ids
|
||||
.iter()
|
||||
.any(|endpoint_id| endpoint_id == &row.endpoint_id)
|
||||
}) && mapping.name == requested_model_name
|
||||
mapping_scope_matches(mapping, row, api_format)
|
||||
&& mapping.name == requested_model_name
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
fn row_has_available_provider_model(
|
||||
row: &StoredMinimalCandidateSelectionRow,
|
||||
api_format: &str,
|
||||
) -> bool {
|
||||
row_mapping_matches_scope(row, api_format)
|
||||
|| row_default_provider_model_name_available(row, api_format)
|
||||
}
|
||||
|
||||
fn row_default_provider_model_name_available(
|
||||
row: &StoredMinimalCandidateSelectionRow,
|
||||
api_format: &str,
|
||||
) -> bool {
|
||||
let Some(mappings) = row.model_provider_model_mappings.as_ref() else {
|
||||
return true;
|
||||
};
|
||||
let mut has_explicit_default_mapping = false;
|
||||
for mapping in mappings {
|
||||
if mapping.name != row.model_provider_model_name {
|
||||
continue;
|
||||
}
|
||||
has_explicit_default_mapping = true;
|
||||
if mapping_scope_matches(mapping, row, api_format) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
!has_explicit_default_mapping
|
||||
}
|
||||
|
||||
fn row_mapping_matches_scope(row: &StoredMinimalCandidateSelectionRow, api_format: &str) -> bool {
|
||||
row.model_provider_model_mappings
|
||||
.as_ref()
|
||||
.is_some_and(|mappings| {
|
||||
mappings
|
||||
.iter()
|
||||
.any(|mapping| mapping_scope_matches(mapping, row, api_format))
|
||||
})
|
||||
}
|
||||
|
||||
fn mapping_scope_matches(
|
||||
mapping: &super::StoredProviderModelMapping,
|
||||
row: &StoredMinimalCandidateSelectionRow,
|
||||
api_format: &str,
|
||||
) -> bool {
|
||||
mapping.api_formats.as_ref().is_none_or(|formats| {
|
||||
formats
|
||||
.iter()
|
||||
.any(|value| api_format_matches(value, api_format))
|
||||
}) && mapping.endpoint_ids.as_ref().is_none_or(|endpoint_ids| {
|
||||
endpoint_ids
|
||||
.iter()
|
||||
.any(|endpoint_id| endpoint_id == &row.endpoint_id)
|
||||
})
|
||||
}
|
||||
|
||||
fn key_auth_channel_matches(row: &CandidateSelectionRow, api_format: &str) -> bool {
|
||||
let provider_type = row.row.provider_type.trim().to_ascii_lowercase();
|
||||
let auth_type = row.row.key_auth_type.trim().to_ascii_lowercase();
|
||||
|
||||
@@ -5,7 +5,8 @@ use std::collections::BTreeSet;
|
||||
|
||||
use super::{
|
||||
MinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
|
||||
StoredPoolKeyCandidateOrder, StoredPoolKeyCandidateRowsQuery, StoredProviderModelMapping,
|
||||
StoredPoolKeyCandidateOrder, StoredPoolKeyCandidateRowsByKeyIdsQuery,
|
||||
StoredPoolKeyCandidateRowsQuery, StoredProviderModelMapping,
|
||||
StoredRequestedModelCandidateRowsQuery,
|
||||
};
|
||||
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)
|
||||
}
|
||||
|
||||
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)]
|
||||
pub struct SqlxMinimalCandidateSelectionReadRepository {
|
||||
pool: PgPool,
|
||||
@@ -704,6 +712,51 @@ impl SqlxMinimalCandidateSelectionReadRepository {
|
||||
}
|
||||
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 {
|
||||
@@ -711,13 +764,110 @@ fn requested_model_selection_sql() -> String {
|
||||
.replace(
|
||||
"AND gm.name = $2",
|
||||
r#"AND (
|
||||
gm.name = $2
|
||||
OR m.provider_model_name = $2
|
||||
(
|
||||
gm.name = $2
|
||||
AND (
|
||||
m.provider_model_mappings IS NULL
|
||||
OR jsonb_typeof(m.provider_model_mappings) <> 'array'
|
||||
OR EXISTS (
|
||||
SELECT 1
|
||||
FROM jsonb_array_elements(
|
||||
CASE
|
||||
WHEN jsonb_typeof(m.provider_model_mappings) = 'array'
|
||||
THEN m.provider_model_mappings
|
||||
ELSE '[]'::jsonb
|
||||
END
|
||||
) AS mapping(value)
|
||||
WHERE (
|
||||
mapping.value -> 'api_formats' IS NULL
|
||||
OR jsonb_typeof(mapping.value -> 'api_formats') <> 'array'
|
||||
OR EXISTS (
|
||||
SELECT 1
|
||||
FROM jsonb_array_elements_text(mapping.value -> 'api_formats') AS fmt(value)
|
||||
WHERE LOWER(BTRIM(fmt.value)) = ANY($3::text[])
|
||||
)
|
||||
)
|
||||
AND (
|
||||
mapping.value -> 'endpoint_ids' IS NULL
|
||||
OR jsonb_typeof(mapping.value -> 'endpoint_ids') <> 'array'
|
||||
OR EXISTS (
|
||||
SELECT 1
|
||||
FROM jsonb_array_elements_text(mapping.value -> 'endpoint_ids') AS endpoint(value)
|
||||
WHERE endpoint.value = pe.id
|
||||
)
|
||||
)
|
||||
)
|
||||
OR NOT EXISTS (
|
||||
SELECT 1
|
||||
FROM jsonb_array_elements(
|
||||
CASE
|
||||
WHEN jsonb_typeof(m.provider_model_mappings) = 'array'
|
||||
THEN m.provider_model_mappings
|
||||
ELSE '[]'::jsonb
|
||||
END
|
||||
) AS mapping(value)
|
||||
WHERE mapping.value ->> 'name' = m.provider_model_name
|
||||
)
|
||||
)
|
||||
)
|
||||
OR (
|
||||
m.provider_model_name = $2
|
||||
AND (
|
||||
m.provider_model_mappings IS NULL
|
||||
OR jsonb_typeof(m.provider_model_mappings) <> 'array'
|
||||
OR NOT EXISTS (
|
||||
SELECT 1
|
||||
FROM jsonb_array_elements(
|
||||
CASE
|
||||
WHEN jsonb_typeof(m.provider_model_mappings) = 'array'
|
||||
THEN m.provider_model_mappings
|
||||
ELSE '[]'::jsonb
|
||||
END
|
||||
) AS mapping(value)
|
||||
WHERE mapping.value ->> 'name' = m.provider_model_name
|
||||
)
|
||||
OR EXISTS (
|
||||
SELECT 1
|
||||
FROM jsonb_array_elements(
|
||||
CASE
|
||||
WHEN jsonb_typeof(m.provider_model_mappings) = 'array'
|
||||
THEN m.provider_model_mappings
|
||||
ELSE '[]'::jsonb
|
||||
END
|
||||
) AS mapping(value)
|
||||
WHERE mapping.value ->> 'name' = m.provider_model_name
|
||||
AND (
|
||||
mapping.value -> 'api_formats' IS NULL
|
||||
OR jsonb_typeof(mapping.value -> 'api_formats') <> 'array'
|
||||
OR EXISTS (
|
||||
SELECT 1
|
||||
FROM jsonb_array_elements_text(mapping.value -> 'api_formats') AS fmt(value)
|
||||
WHERE LOWER(BTRIM(fmt.value)) = ANY($3::text[])
|
||||
)
|
||||
)
|
||||
AND (
|
||||
mapping.value -> 'endpoint_ids' IS NULL
|
||||
OR jsonb_typeof(mapping.value -> 'endpoint_ids') <> 'array'
|
||||
OR EXISTS (
|
||||
SELECT 1
|
||||
FROM jsonb_array_elements_text(mapping.value -> 'endpoint_ids') AS endpoint(value)
|
||||
WHERE endpoint.value = pe.id
|
||||
)
|
||||
)
|
||||
)
|
||||
)
|
||||
)
|
||||
OR (
|
||||
jsonb_typeof(m.provider_model_mappings) = 'array'
|
||||
AND EXISTS (
|
||||
SELECT 1
|
||||
FROM jsonb_array_elements(m.provider_model_mappings) AS mapping(value)
|
||||
FROM jsonb_array_elements(
|
||||
CASE
|
||||
WHEN jsonb_typeof(m.provider_model_mappings) = 'array'
|
||||
THEN m.provider_model_mappings
|
||||
ELSE '[]'::jsonb
|
||||
END
|
||||
) AS mapping(value)
|
||||
WHERE mapping.value ->> 'name' = $2
|
||||
AND (
|
||||
mapping.value -> 'api_formats' IS NULL
|
||||
@@ -820,6 +970,13 @@ impl MinimalCandidateSelectionReadRepository for SqlxMinimalCandidateSelectionRe
|
||||
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
|
||||
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(
|
||||
@@ -1108,7 +1265,7 @@ mod tests {
|
||||
let sql = requested_model_selection_sql();
|
||||
|
||||
assert!(sql.contains("m.provider_model_name = $2"));
|
||||
assert!(sql.contains("jsonb_array_elements(m.provider_model_mappings)"));
|
||||
assert!(sql.contains("jsonb_array_elements("));
|
||||
assert!(!sql.contains("json_typeof(m.provider_model_mappings)"));
|
||||
assert!(!sql.contains("json_array_elements_text(gm.config -> 'model_mappings')"));
|
||||
assert!(sql.contains("ORDER BY\n global_model_name ASC,"));
|
||||
|
||||
@@ -5,7 +5,8 @@ use sqlx::{sqlite::SqliteRow, QueryBuilder, Row, Sqlite};
|
||||
|
||||
use super::{
|
||||
MinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
|
||||
StoredPoolKeyCandidateOrder, StoredPoolKeyCandidateRowsQuery, StoredProviderModelMapping,
|
||||
StoredPoolKeyCandidateOrder, StoredPoolKeyCandidateRowsByKeyIdsQuery,
|
||||
StoredPoolKeyCandidateRowsQuery, StoredProviderModelMapping,
|
||||
StoredRequestedModelCandidateRowsQuery,
|
||||
};
|
||||
use crate::driver::sqlite::SqlitePool;
|
||||
@@ -196,6 +197,40 @@ impl MinimalCandidateSelectionReadRepository for SqliteMinimalCandidateSelection
|
||||
.map(|item| item.row)
|
||||
.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> {
|
||||
@@ -313,26 +348,75 @@ fn row_matches_requested_model(
|
||||
requested_model_name: &str,
|
||||
api_format: &str,
|
||||
) -> bool {
|
||||
row.global_model_name == requested_model_name
|
||||
|| row.model_provider_model_name == requested_model_name
|
||||
(row_has_available_provider_model(row, api_format)
|
||||
&& row.global_model_name == requested_model_name)
|
||||
|| (row_default_provider_model_name_available(row, api_format)
|
||||
&& row.model_provider_model_name == requested_model_name)
|
||||
|| row
|
||||
.model_provider_model_mappings
|
||||
.as_ref()
|
||||
.is_some_and(|mappings| {
|
||||
mappings.iter().any(|mapping| {
|
||||
mapping.api_formats.as_ref().is_none_or(|formats| {
|
||||
formats
|
||||
.iter()
|
||||
.any(|value| api_format_matches(value, api_format))
|
||||
}) && mapping.endpoint_ids.as_ref().is_none_or(|endpoint_ids| {
|
||||
endpoint_ids
|
||||
.iter()
|
||||
.any(|endpoint_id| endpoint_id == &row.endpoint_id)
|
||||
}) && mapping.name == requested_model_name
|
||||
mapping_scope_matches(mapping, row, api_format)
|
||||
&& mapping.name == requested_model_name
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
fn row_has_available_provider_model(
|
||||
row: &StoredMinimalCandidateSelectionRow,
|
||||
api_format: &str,
|
||||
) -> bool {
|
||||
row_mapping_matches_scope(row, api_format)
|
||||
|| row_default_provider_model_name_available(row, api_format)
|
||||
}
|
||||
|
||||
fn row_default_provider_model_name_available(
|
||||
row: &StoredMinimalCandidateSelectionRow,
|
||||
api_format: &str,
|
||||
) -> bool {
|
||||
let Some(mappings) = row.model_provider_model_mappings.as_ref() else {
|
||||
return true;
|
||||
};
|
||||
let mut has_explicit_default_mapping = false;
|
||||
for mapping in mappings {
|
||||
if mapping.name != row.model_provider_model_name {
|
||||
continue;
|
||||
}
|
||||
has_explicit_default_mapping = true;
|
||||
if mapping_scope_matches(mapping, row, api_format) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
!has_explicit_default_mapping
|
||||
}
|
||||
|
||||
fn row_mapping_matches_scope(row: &StoredMinimalCandidateSelectionRow, api_format: &str) -> bool {
|
||||
row.model_provider_model_mappings
|
||||
.as_ref()
|
||||
.is_some_and(|mappings| {
|
||||
mappings
|
||||
.iter()
|
||||
.any(|mapping| mapping_scope_matches(mapping, row, api_format))
|
||||
})
|
||||
}
|
||||
|
||||
fn mapping_scope_matches(
|
||||
mapping: &super::StoredProviderModelMapping,
|
||||
row: &StoredMinimalCandidateSelectionRow,
|
||||
api_format: &str,
|
||||
) -> bool {
|
||||
mapping.api_formats.as_ref().is_none_or(|formats| {
|
||||
formats
|
||||
.iter()
|
||||
.any(|value| api_format_matches(value, api_format))
|
||||
}) && mapping.endpoint_ids.as_ref().is_none_or(|endpoint_ids| {
|
||||
endpoint_ids
|
||||
.iter()
|
||||
.any(|endpoint_id| endpoint_id == &row.endpoint_id)
|
||||
})
|
||||
}
|
||||
|
||||
fn key_auth_channel_matches(row: &CandidateSelectionRow, api_format: &str) -> bool {
|
||||
let provider_type = row.row.provider_type.trim().to_ascii_lowercase();
|
||||
let auth_type = row.row.key_auth_type.trim().to_ascii_lowercase();
|
||||
|
||||
@@ -16,6 +16,7 @@ pub mod gemini_file_mappings;
|
||||
pub mod global_models;
|
||||
pub mod management_tokens;
|
||||
pub mod oauth_providers;
|
||||
pub mod pool_scores;
|
||||
pub mod provider_catalog;
|
||||
pub mod provider_oauth;
|
||||
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")))
|
||||
}
|
||||
16
crates/aether-dispatch-core/Cargo.toml
Normal file
16
crates/aether-dispatch-core/Cargo.toml
Normal file
@@ -0,0 +1,16 @@
|
||||
[package]
|
||||
name = "aether-dispatch-core"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
description = "Request-scoped dispatch sequence and pool cursor primitives for Aether"
|
||||
|
||||
[dependencies]
|
||||
async-trait.workspace = true
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
thiserror.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
tokio.workspace = true
|
||||
68
crates/aether-dispatch-core/src/candidate.rs
Normal file
68
crates/aether-dispatch-core/src/candidate.rs
Normal file
@@ -0,0 +1,68 @@
|
||||
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct ProviderEndpointRef {
|
||||
pub provider_id: String,
|
||||
pub endpoint_id: String,
|
||||
pub model_id: String,
|
||||
pub selected_provider_model_name: String,
|
||||
pub api_format: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct KeyRef {
|
||||
pub provider_id: String,
|
||||
pub endpoint_id: String,
|
||||
pub key_id: String,
|
||||
pub model_id: String,
|
||||
pub selected_provider_model_name: String,
|
||||
pub api_format: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct PoolRef {
|
||||
pub provider_id: String,
|
||||
pub endpoint_id: String,
|
||||
pub model_id: String,
|
||||
pub selected_provider_model_name: String,
|
||||
pub api_format: String,
|
||||
pub pool_group_id: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct DispatchRankFacts {
|
||||
pub provider_priority: i32,
|
||||
pub key_priority: Option<i32>,
|
||||
pub ranking_reason: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
pub enum DispatchCandidateRef {
|
||||
SingleKey {
|
||||
key: KeyRef,
|
||||
rank: DispatchRankFacts,
|
||||
},
|
||||
PoolRef {
|
||||
pool: PoolRef,
|
||||
rank: DispatchRankFacts,
|
||||
},
|
||||
}
|
||||
|
||||
impl DispatchCandidateRef {
|
||||
pub fn provider_endpoint(&self) -> ProviderEndpointRef {
|
||||
match self {
|
||||
Self::SingleKey { key, .. } => ProviderEndpointRef {
|
||||
provider_id: key.provider_id.clone(),
|
||||
endpoint_id: key.endpoint_id.clone(),
|
||||
model_id: key.model_id.clone(),
|
||||
selected_provider_model_name: key.selected_provider_model_name.clone(),
|
||||
api_format: key.api_format.clone(),
|
||||
},
|
||||
Self::PoolRef { pool, .. } => ProviderEndpointRef {
|
||||
provider_id: pool.provider_id.clone(),
|
||||
endpoint_id: pool.endpoint_id.clone(),
|
||||
model_id: pool.model_id.clone(),
|
||||
selected_provider_model_name: pool.selected_provider_model_name.clone(),
|
||||
api_format: pool.api_format.clone(),
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
16
crates/aether-dispatch-core/src/effects.rs
Normal file
16
crates/aether-dispatch-core/src/effects.rs
Normal file
@@ -0,0 +1,16 @@
|
||||
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
pub enum DispatchEffectKind {
|
||||
CandidateFailed,
|
||||
RateLimited,
|
||||
Succeeded,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct DispatchEffect {
|
||||
pub kind: DispatchEffectKind,
|
||||
pub provider_id: String,
|
||||
pub endpoint_id: String,
|
||||
pub key_id: Option<String>,
|
||||
pub candidate_index: u32,
|
||||
pub reason: Option<String>,
|
||||
}
|
||||
15
crates/aether-dispatch-core/src/lib.rs
Normal file
15
crates/aether-dispatch-core/src/lib.rs
Normal file
@@ -0,0 +1,15 @@
|
||||
pub mod candidate;
|
||||
pub mod effects;
|
||||
pub mod pool;
|
||||
pub mod sequence;
|
||||
|
||||
pub use candidate::{
|
||||
DispatchCandidateRef, DispatchRankFacts, KeyRef, PoolRef, ProviderEndpointRef,
|
||||
};
|
||||
pub use effects::{DispatchEffect, DispatchEffectKind};
|
||||
pub use pool::{
|
||||
run_pool_dispatch_cursor, PoolDispatchCursorOutcome, PoolDispatchError, PoolDispatchPort,
|
||||
PoolDispatchWindow, PoolWindowConfig, DEFAULT_POOL_MAX_SCAN, DEFAULT_POOL_PAGE_SIZE,
|
||||
DEFAULT_POOL_WINDOW_SIZE,
|
||||
};
|
||||
pub use sequence::{DispatchSequence, DispatchSequenceItem, DispatchSequenceMark};
|
||||
216
crates/aether-dispatch-core/src/pool.rs
Normal file
216
crates/aether-dispatch-core/src/pool.rs
Normal file
@@ -0,0 +1,216 @@
|
||||
use async_trait::async_trait;
|
||||
|
||||
pub const DEFAULT_POOL_WINDOW_SIZE: u32 = 16;
|
||||
pub const DEFAULT_POOL_PAGE_SIZE: u32 = 64;
|
||||
pub const DEFAULT_POOL_MAX_SCAN: u32 = 512;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct PoolWindowConfig {
|
||||
pub window_size: u32,
|
||||
pub page_size: u32,
|
||||
pub max_scan: u32,
|
||||
}
|
||||
|
||||
impl Default for PoolWindowConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
window_size: DEFAULT_POOL_WINDOW_SIZE,
|
||||
page_size: DEFAULT_POOL_PAGE_SIZE,
|
||||
max_scan: DEFAULT_POOL_MAX_SCAN,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl PoolWindowConfig {
|
||||
pub fn normalized(self) -> Self {
|
||||
let page_size = self.page_size.max(1);
|
||||
let window_size = self.window_size.max(1).min(page_size);
|
||||
let max_scan = self.max_scan.max(window_size);
|
||||
Self {
|
||||
window_size,
|
||||
page_size,
|
||||
max_scan,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct PoolDispatchWindow<Candidate> {
|
||||
pub candidates: Vec<Candidate>,
|
||||
pub scanned_count: u32,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct PoolDispatchCursorOutcome<Candidate> {
|
||||
pub candidates: Vec<Candidate>,
|
||||
pub scanned_count: u32,
|
||||
pub exhausted: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
|
||||
pub enum PoolDispatchError<Error> {
|
||||
#[error("pool dispatch port failed")]
|
||||
Port(Error),
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait PoolDispatchPort {
|
||||
type Candidate: Send;
|
||||
type Error: Send;
|
||||
|
||||
async fn read_page(
|
||||
&mut self,
|
||||
offset: u32,
|
||||
limit: u32,
|
||||
) -> Result<Vec<Self::Candidate>, Self::Error>;
|
||||
|
||||
async fn rank_and_filter_window(
|
||||
&mut self,
|
||||
candidates: Vec<Self::Candidate>,
|
||||
window_size: u32,
|
||||
) -> Result<PoolDispatchWindow<Self::Candidate>, Self::Error>;
|
||||
}
|
||||
|
||||
pub async fn run_pool_dispatch_cursor<Port>(
|
||||
port: &mut Port,
|
||||
config: PoolWindowConfig,
|
||||
) -> Result<PoolDispatchCursorOutcome<Port::Candidate>, PoolDispatchError<Port::Error>>
|
||||
where
|
||||
Port: PoolDispatchPort + Send,
|
||||
{
|
||||
let config = config.normalized();
|
||||
let mut offset = 0_u32;
|
||||
let mut scanned_count = 0_u32;
|
||||
|
||||
while scanned_count < config.max_scan {
|
||||
let limit = config.page_size.min(config.max_scan - scanned_count);
|
||||
let page = port
|
||||
.read_page(offset, limit)
|
||||
.await
|
||||
.map_err(PoolDispatchError::Port)?;
|
||||
if page.is_empty() {
|
||||
return Ok(PoolDispatchCursorOutcome {
|
||||
candidates: Vec::new(),
|
||||
scanned_count,
|
||||
exhausted: true,
|
||||
});
|
||||
}
|
||||
|
||||
let page_len = u32::try_from(page.len()).unwrap_or(u32::MAX);
|
||||
offset = offset.saturating_add(page_len);
|
||||
scanned_count = scanned_count.saturating_add(page_len);
|
||||
|
||||
let window = port
|
||||
.rank_and_filter_window(page, config.window_size)
|
||||
.await
|
||||
.map_err(PoolDispatchError::Port)?;
|
||||
if !window.candidates.is_empty() {
|
||||
return Ok(PoolDispatchCursorOutcome {
|
||||
candidates: window.candidates,
|
||||
scanned_count,
|
||||
exhausted: false,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
Ok(PoolDispatchCursorOutcome {
|
||||
candidates: Vec::new(),
|
||||
scanned_count,
|
||||
exhausted: true,
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::VecDeque;
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
use super::{run_pool_dispatch_cursor, PoolDispatchPort, PoolDispatchWindow, PoolWindowConfig};
|
||||
|
||||
#[derive(Default)]
|
||||
struct TestPort {
|
||||
pages: VecDeque<Vec<u32>>,
|
||||
read_limits: Vec<u32>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl PoolDispatchPort for TestPort {
|
||||
type Candidate = u32;
|
||||
type Error = std::convert::Infallible;
|
||||
|
||||
async fn read_page(
|
||||
&mut self,
|
||||
_offset: u32,
|
||||
limit: u32,
|
||||
) -> Result<Vec<Self::Candidate>, Self::Error> {
|
||||
self.read_limits.push(limit);
|
||||
Ok(self.pages.pop_front().unwrap_or_default())
|
||||
}
|
||||
|
||||
async fn rank_and_filter_window(
|
||||
&mut self,
|
||||
mut candidates: Vec<Self::Candidate>,
|
||||
window_size: u32,
|
||||
) -> Result<PoolDispatchWindow<Self::Candidate>, Self::Error> {
|
||||
candidates.sort();
|
||||
candidates.truncate(window_size as usize);
|
||||
let scanned_count = u32::try_from(candidates.len()).unwrap_or(u32::MAX);
|
||||
Ok(PoolDispatchWindow {
|
||||
candidates,
|
||||
scanned_count,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn small_pool_is_returned_in_one_frozen_window() {
|
||||
let mut port = TestPort {
|
||||
pages: VecDeque::from([vec![3, 1, 2]]),
|
||||
read_limits: Vec::new(),
|
||||
};
|
||||
|
||||
let outcome = run_pool_dispatch_cursor(&mut port, PoolWindowConfig::default())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(outcome.candidates, [1, 2, 3]);
|
||||
assert_eq!(outcome.scanned_count, 3);
|
||||
assert_eq!(port.read_limits, [64]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn large_pool_returns_bounded_window() {
|
||||
let mut port = TestPort {
|
||||
pages: VecDeque::from([(0..100).rev().collect::<Vec<_>>()]),
|
||||
read_limits: Vec::new(),
|
||||
};
|
||||
|
||||
let outcome = run_pool_dispatch_cursor(&mut port, PoolWindowConfig::default())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(outcome.candidates.len(), 16);
|
||||
assert_eq!(outcome.candidates[0], 0);
|
||||
assert_eq!(outcome.candidates[15], 15);
|
||||
assert_eq!(port.read_limits, [64]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn max_scan_caps_page_reads() {
|
||||
let mut port = TestPort {
|
||||
pages: VecDeque::from([Vec::new()]),
|
||||
read_limits: Vec::new(),
|
||||
};
|
||||
let config = PoolWindowConfig {
|
||||
window_size: 16,
|
||||
page_size: 64,
|
||||
max_scan: 32,
|
||||
};
|
||||
|
||||
let outcome = run_pool_dispatch_cursor(&mut port, config).await.unwrap();
|
||||
|
||||
assert!(outcome.exhausted);
|
||||
assert_eq!(port.read_limits, [32]);
|
||||
}
|
||||
}
|
||||
110
crates/aether-dispatch-core/src/sequence.rs
Normal file
110
crates/aether-dispatch-core/src/sequence.rs
Normal file
@@ -0,0 +1,110 @@
|
||||
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct DispatchSequenceItem<Candidate> {
|
||||
pub candidate_index: u32,
|
||||
pub retry_index: u32,
|
||||
pub candidate: Candidate,
|
||||
pub mark: DispatchSequenceMark,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
pub enum DispatchSequenceMark {
|
||||
Pending,
|
||||
Failed,
|
||||
Succeeded,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct DispatchSequence<Candidate> {
|
||||
items: Vec<DispatchSequenceItem<Candidate>>,
|
||||
cursor: usize,
|
||||
}
|
||||
|
||||
impl<Candidate> DispatchSequence<Candidate> {
|
||||
pub fn new(items: Vec<DispatchSequenceItem<Candidate>>) -> Self {
|
||||
Self { items, cursor: 0 }
|
||||
}
|
||||
|
||||
pub fn from_candidates(candidates: Vec<Candidate>) -> Self {
|
||||
let items = candidates
|
||||
.into_iter()
|
||||
.enumerate()
|
||||
.map(|(index, candidate)| DispatchSequenceItem {
|
||||
candidate_index: u32::try_from(index).unwrap_or(u32::MAX),
|
||||
retry_index: 0,
|
||||
candidate,
|
||||
mark: DispatchSequenceMark::Pending,
|
||||
})
|
||||
.collect();
|
||||
Self::new(items)
|
||||
}
|
||||
|
||||
pub fn peek_current(&self) -> Option<&DispatchSequenceItem<Candidate>> {
|
||||
self.items.get(self.cursor)
|
||||
}
|
||||
|
||||
#[allow(clippy::should_implement_trait)]
|
||||
pub fn next(&mut self) -> Option<&DispatchSequenceItem<Candidate>> {
|
||||
while self
|
||||
.items
|
||||
.get(self.cursor)
|
||||
.is_some_and(|item| item.mark != DispatchSequenceMark::Pending)
|
||||
{
|
||||
self.cursor = self.cursor.saturating_add(1);
|
||||
}
|
||||
self.items.get(self.cursor)
|
||||
}
|
||||
|
||||
pub fn mark_failed(&mut self) -> Option<&DispatchSequenceItem<Candidate>> {
|
||||
self.mark_current(DispatchSequenceMark::Failed)
|
||||
}
|
||||
|
||||
pub fn mark_succeeded(&mut self) -> Option<&DispatchSequenceItem<Candidate>> {
|
||||
self.mark_current(DispatchSequenceMark::Succeeded)
|
||||
}
|
||||
|
||||
pub fn len(&self) -> usize {
|
||||
self.items.len()
|
||||
}
|
||||
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.items.is_empty()
|
||||
}
|
||||
|
||||
pub fn cursor(&self) -> usize {
|
||||
self.cursor
|
||||
}
|
||||
|
||||
pub fn items(&self) -> &[DispatchSequenceItem<Candidate>] {
|
||||
&self.items
|
||||
}
|
||||
|
||||
fn mark_current(
|
||||
&mut self,
|
||||
mark: DispatchSequenceMark,
|
||||
) -> Option<&DispatchSequenceItem<Candidate>> {
|
||||
let item = self.items.get_mut(self.cursor)?;
|
||||
item.mark = mark;
|
||||
self.cursor = self.cursor.saturating_add(1);
|
||||
self.items.get(self.cursor)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{DispatchSequence, DispatchSequenceMark};
|
||||
|
||||
#[test]
|
||||
fn mark_failed_advances_without_reordering() {
|
||||
let mut sequence = DispatchSequence::from_candidates(vec!["a", "b", "c"]);
|
||||
|
||||
assert_eq!(sequence.next().map(|item| item.candidate), Some("a"));
|
||||
assert_eq!(sequence.mark_failed().map(|item| item.candidate), Some("b"));
|
||||
assert_eq!(sequence.next().map(|item| item.candidate), Some("b"));
|
||||
assert_eq!(sequence.mark_failed().map(|item| item.candidate), Some("c"));
|
||||
assert_eq!(sequence.next().map(|item| item.candidate), Some("c"));
|
||||
|
||||
assert_eq!(sequence.items()[0].mark, DispatchSequenceMark::Failed);
|
||||
assert_eq!(sequence.items()[1].mark, DispatchSequenceMark::Failed);
|
||||
assert_eq!(sequence.items()[2].mark, DispatchSequenceMark::Pending);
|
||||
}
|
||||
}
|
||||
@@ -29,33 +29,38 @@ pub fn resolve_requested_global_model_name_with_model_directives(
|
||||
requested_model_name_candidates(requested_model_name, enable_model_directives).find_map(
|
||||
|requested_model_name| {
|
||||
let requested_model_name = requested_model_name.as_ref();
|
||||
resolve_global_model_name_by(rows, |row| row.global_model_name == requested_model_name)
|
||||
.or_else(|| {
|
||||
resolve_global_model_name_by(rows, |row| {
|
||||
row.model_provider_model_name == requested_model_name
|
||||
})
|
||||
resolve_global_model_name_by(rows, |row| {
|
||||
row_has_available_provider_model(row, api_format)
|
||||
&& row.global_model_name == requested_model_name
|
||||
})
|
||||
.or_else(|| {
|
||||
resolve_global_model_name_by(rows, |row| {
|
||||
row_default_provider_model_name_available(row, api_format)
|
||||
&& row.model_provider_model_name == requested_model_name
|
||||
})
|
||||
.or_else(|| {
|
||||
resolve_global_model_name_by(rows, |row| {
|
||||
row.model_provider_model_mappings
|
||||
.as_ref()
|
||||
.is_some_and(|mappings| {
|
||||
mappings.iter().any(|mapping| {
|
||||
mapping_scope_matches(mapping, row, api_format)
|
||||
&& mapping.name == requested_model_name
|
||||
})
|
||||
})
|
||||
.or_else(|| {
|
||||
resolve_global_model_name_by(rows, |row| {
|
||||
row.model_provider_model_mappings
|
||||
.as_ref()
|
||||
.is_some_and(|mappings| {
|
||||
mappings.iter().any(|mapping| {
|
||||
mapping_scope_matches(mapping, row, api_format)
|
||||
&& mapping.name == requested_model_name
|
||||
})
|
||||
})
|
||||
})
|
||||
})
|
||||
.or_else(|| {
|
||||
resolve_global_model_name_by(rows, |row| {
|
||||
row.global_model_mappings.as_ref().is_some_and(|patterns| {
|
||||
})
|
||||
.or_else(|| {
|
||||
resolve_global_model_name_by(rows, |row| {
|
||||
row_has_available_provider_model(row, api_format)
|
||||
&& row.global_model_mappings.as_ref().is_some_and(|patterns| {
|
||||
patterns
|
||||
.iter()
|
||||
.any(|pattern| matches_model_mapping(pattern, requested_model_name))
|
||||
})
|
||||
})
|
||||
})
|
||||
})
|
||||
},
|
||||
)
|
||||
}
|
||||
@@ -86,8 +91,15 @@ fn row_supports_requested_model_exact(
|
||||
requested_model_name: &str,
|
||||
api_format: &str,
|
||||
) -> bool {
|
||||
row.global_model_name == requested_model_name
|
||||
|| row.model_provider_model_name == requested_model_name
|
||||
row_has_available_provider_model(row, api_format)
|
||||
&& (row.global_model_name == requested_model_name
|
||||
|| (row_default_provider_model_name_available(row, api_format)
|
||||
&& row.model_provider_model_name == requested_model_name)
|
||||
|| row.global_model_mappings.as_ref().is_some_and(|patterns| {
|
||||
patterns
|
||||
.iter()
|
||||
.any(|pattern| matches_model_mapping(pattern, requested_model_name))
|
||||
}))
|
||||
|| row
|
||||
.model_provider_model_mappings
|
||||
.as_ref()
|
||||
@@ -97,11 +109,6 @@ fn row_supports_requested_model_exact(
|
||||
&& mapping.name == requested_model_name
|
||||
})
|
||||
})
|
||||
|| row.global_model_mappings.as_ref().is_some_and(|patterns| {
|
||||
patterns
|
||||
.iter()
|
||||
.any(|pattern| matches_model_mapping(pattern, requested_model_name))
|
||||
})
|
||||
}
|
||||
|
||||
fn resolve_global_model_name_by<F>(
|
||||
@@ -138,7 +145,7 @@ pub fn resolve_provider_model_name_with_model_directives(
|
||||
api_format: &str,
|
||||
enable_model_directives: bool,
|
||||
) -> Option<(String, Option<String>)> {
|
||||
let selected_provider_model_name = select_provider_model_name(row, api_format);
|
||||
let selected_provider_model_name = resolve_selected_provider_model_name(row, api_format)?;
|
||||
let Some(key_allowed_models) = row.key_allowed_models.as_ref() else {
|
||||
return Some((selected_provider_model_name, None));
|
||||
};
|
||||
@@ -195,11 +202,19 @@ pub fn select_provider_model_name(
|
||||
row: &StoredMinimalCandidateSelectionRow,
|
||||
api_format: &str,
|
||||
) -> String {
|
||||
resolve_selected_provider_model_name(row, api_format)
|
||||
.unwrap_or_else(|| row.model_provider_model_name.clone())
|
||||
}
|
||||
|
||||
fn resolve_selected_provider_model_name(
|
||||
row: &StoredMinimalCandidateSelectionRow,
|
||||
api_format: &str,
|
||||
) -> Option<String> {
|
||||
let Some(mappings) = row.model_provider_model_mappings.as_ref() else {
|
||||
return row.model_provider_model_name.clone();
|
||||
return Some(row.model_provider_model_name.clone());
|
||||
};
|
||||
|
||||
mappings
|
||||
if let Some(mapping) = mappings
|
||||
.iter()
|
||||
.filter(|mapping| mapping_scope_matches(mapping, row, api_format))
|
||||
.min_by(|left, right| {
|
||||
@@ -207,15 +222,22 @@ pub fn select_provider_model_name(
|
||||
.cmp(&right.priority)
|
||||
.then(left.name.cmp(&right.name))
|
||||
})
|
||||
.map(|mapping| mapping.name.clone())
|
||||
.unwrap_or_else(|| row.model_provider_model_name.clone())
|
||||
{
|
||||
return Some(mapping.name.clone());
|
||||
}
|
||||
|
||||
row_default_provider_model_name_available(row, api_format)
|
||||
.then(|| row.model_provider_model_name.clone())
|
||||
}
|
||||
|
||||
pub fn candidate_model_names(
|
||||
row: &StoredMinimalCandidateSelectionRow,
|
||||
api_format: &str,
|
||||
) -> BTreeSet<String> {
|
||||
let mut names = BTreeSet::from([row.model_provider_model_name.clone()]);
|
||||
let mut names = BTreeSet::new();
|
||||
if row_default_provider_model_name_available(row, api_format) {
|
||||
names.insert(row.model_provider_model_name.clone());
|
||||
}
|
||||
if let Some(mappings) = row.model_provider_model_mappings.as_ref() {
|
||||
for mapping in mappings {
|
||||
if mapping_scope_matches(mapping, row, api_format) {
|
||||
@@ -226,6 +248,33 @@ pub fn candidate_model_names(
|
||||
names
|
||||
}
|
||||
|
||||
fn row_has_available_provider_model(
|
||||
row: &StoredMinimalCandidateSelectionRow,
|
||||
api_format: &str,
|
||||
) -> bool {
|
||||
resolve_selected_provider_model_name(row, api_format).is_some()
|
||||
}
|
||||
|
||||
fn row_default_provider_model_name_available(
|
||||
row: &StoredMinimalCandidateSelectionRow,
|
||||
api_format: &str,
|
||||
) -> bool {
|
||||
let Some(mappings) = row.model_provider_model_mappings.as_ref() else {
|
||||
return true;
|
||||
};
|
||||
let mut has_explicit_default_mapping = false;
|
||||
for mapping in mappings {
|
||||
if mapping.name != row.model_provider_model_name {
|
||||
continue;
|
||||
}
|
||||
has_explicit_default_mapping = true;
|
||||
if mapping_scope_matches(mapping, row, api_format) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
!has_explicit_default_mapping
|
||||
}
|
||||
|
||||
fn mapping_scope_matches(
|
||||
mapping: &StoredProviderModelMapping,
|
||||
row: &StoredMinimalCandidateSelectionRow,
|
||||
@@ -358,7 +407,8 @@ fn row_has_candidate_model_name(
|
||||
api_format: &str,
|
||||
model_name: &str,
|
||||
) -> bool {
|
||||
row.model_provider_model_name == model_name
|
||||
(row_default_provider_model_name_available(row, api_format)
|
||||
&& row.model_provider_model_name == model_name)
|
||||
|| row
|
||||
.model_provider_model_mappings
|
||||
.as_ref()
|
||||
@@ -494,6 +544,54 @@ mod tests {
|
||||
assert_eq!(resolved.1.as_deref(), Some("gpt-5.4"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn endpoint_scoped_default_mapping_limits_exact_global_model_match() {
|
||||
let mut row = sample_row("deepseek-v4-pro", "deepseek-v4-pro");
|
||||
row.endpoint_id = "endpoint-claude".to_string();
|
||||
row.endpoint_api_format = "claude:messages".to_string();
|
||||
row.model_provider_model_mappings = Some(vec![StoredProviderModelMapping {
|
||||
name: "deepseek-v4-pro".to_string(),
|
||||
priority: 1,
|
||||
api_formats: None,
|
||||
endpoint_ids: Some(vec!["endpoint-openai".to_string()]),
|
||||
}]);
|
||||
|
||||
assert!(!row_supports_requested_model(
|
||||
&row,
|
||||
"deepseek-v4-pro",
|
||||
"claude:messages"
|
||||
));
|
||||
assert!(resolve_provider_model_name(&row, "deepseek-v4-pro", "claude:messages").is_none());
|
||||
assert_eq!(
|
||||
resolve_requested_global_model_name_with_model_directives(
|
||||
&[row.clone()],
|
||||
"deepseek-v4-pro",
|
||||
"claude:messages",
|
||||
false,
|
||||
),
|
||||
None
|
||||
);
|
||||
|
||||
row.endpoint_id = "endpoint-openai".to_string();
|
||||
row.endpoint_api_format = "openai:chat".to_string();
|
||||
|
||||
assert!(row_supports_requested_model(
|
||||
&row,
|
||||
"deepseek-v4-pro",
|
||||
"openai:chat"
|
||||
));
|
||||
assert_eq!(
|
||||
resolve_requested_global_model_name_with_model_directives(
|
||||
&[row],
|
||||
"deepseek-v4-pro",
|
||||
"openai:chat",
|
||||
false,
|
||||
)
|
||||
.as_deref(),
|
||||
Some("deepseek-v4-pro")
|
||||
);
|
||||
}
|
||||
|
||||
fn sample_row(
|
||||
global_model_name: &str,
|
||||
model_provider_model_name: &str,
|
||||
|
||||
Reference in New Issue
Block a user