mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-03 01:40:21 +08:00
fix(build): 修复网关构建失败并收口候选选择与 finalize 回归
- 补齐 DecisionTraceCandidate 新增字段,修复审计测试构造 - 修正 usage 内存仓库的 created_at_unix_ms 字段引用与秒/毫秒换算 - 将 build_minimal_candidate_selection 重构为输入对象,消除 clippy 参数过多问题 - 修复 admin global model created_at 旧字段残留引用 - 修复 openai:cli 与 openai:compact 同家族 finalize 在 needs_conversion=true 时的成功回落逻辑 - 清理 aether-gateway 中的 derive/default 与 needless borrow 等 clippy 问题
This commit is contained in:
@@ -192,7 +192,7 @@ async fn resolve_tunnel_owner_affinity_from_transport(
|
|||||||
state: PlannerAppState<'_>,
|
state: PlannerAppState<'_>,
|
||||||
transport: &GatewayProviderTransportSnapshot,
|
transport: &GatewayProviderTransportSnapshot,
|
||||||
) -> TunnelOwnerAffinityBucket {
|
) -> TunnelOwnerAffinityBucket {
|
||||||
let Some(proxy) = resolve_transport_proxy_snapshot(&transport) else {
|
let Some(proxy) = resolve_transport_proxy_snapshot(transport) else {
|
||||||
return TunnelOwnerAffinityBucket::Neutral;
|
return TunnelOwnerAffinityBucket::Neutral;
|
||||||
};
|
};
|
||||||
if proxy.enabled == Some(false) {
|
if proxy.enabled == Some(false) {
|
||||||
|
|||||||
@@ -52,8 +52,7 @@ pub(super) async fn build_cross_format_local_openai_chat_decision_payload_for_ca
|
|||||||
else {
|
else {
|
||||||
return None;
|
return None;
|
||||||
};
|
};
|
||||||
if !request_pair_allowed_for_transport(&transport, "openai:chat", provider_api_format.as_str())
|
if !request_pair_allowed_for_transport(transport, "openai:chat", provider_api_format.as_str()) {
|
||||||
{
|
|
||||||
let skip_reason =
|
let skip_reason =
|
||||||
if request_conversion_requires_enable_flag("openai:chat", provider_api_format.as_str())
|
if request_conversion_requires_enable_flag("openai:chat", provider_api_format.as_str())
|
||||||
&& !transport.provider.enable_format_conversion
|
&& !transport.provider.enable_format_conversion
|
||||||
|
|||||||
@@ -3,8 +3,8 @@ use aether_data_contracts::repository::candidate_selection::StoredMinimalCandida
|
|||||||
use aether_scheduler_core::{
|
use aether_scheduler_core::{
|
||||||
auth_constraints_allow_api_format, build_minimal_candidate_selection,
|
auth_constraints_allow_api_format, build_minimal_candidate_selection,
|
||||||
collect_global_model_names_for_required_capability, normalize_api_format,
|
collect_global_model_names_for_required_capability, normalize_api_format,
|
||||||
resolve_requested_global_model_name, SchedulerAuthConstraints,
|
resolve_requested_global_model_name, BuildMinimalCandidateSelectionInput,
|
||||||
SchedulerMinimalCandidateSelectionCandidate, SchedulerPriorityMode,
|
SchedulerAuthConstraints, SchedulerMinimalCandidateSelectionCandidate, SchedulerPriorityMode,
|
||||||
};
|
};
|
||||||
use async_trait::async_trait;
|
use async_trait::async_trait;
|
||||||
use std::collections::BTreeSet;
|
use std::collections::BTreeSet;
|
||||||
@@ -167,17 +167,17 @@ pub(crate) async fn read_minimal_candidate_selection_with_priority_mode_and_affi
|
|||||||
return Ok(Vec::new());
|
return Ok(Vec::new());
|
||||||
};
|
};
|
||||||
let auth_constraints = auth_snapshot.map(auth_snapshot_constraints);
|
let auth_constraints = auth_snapshot.map(auth_snapshot_constraints);
|
||||||
build_minimal_candidate_selection(
|
build_minimal_candidate_selection(BuildMinimalCandidateSelectionInput {
|
||||||
rows,
|
rows,
|
||||||
&normalized_api_format,
|
normalized_api_format: &normalized_api_format,
|
||||||
requested_model_name,
|
requested_model_name,
|
||||||
resolved_global_model_name.as_str(),
|
resolved_global_model_name: resolved_global_model_name.as_str(),
|
||||||
require_streaming,
|
require_streaming,
|
||||||
required_capabilities,
|
required_capabilities,
|
||||||
auth_constraints.as_ref(),
|
auth_constraints: auth_constraints.as_ref(),
|
||||||
affinity_key,
|
affinity_key,
|
||||||
priority_mode,
|
priority_mode,
|
||||||
)
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(crate) async fn read_global_model_names_for_required_capability(
|
pub(crate) async fn read_global_model_names_for_required_capability(
|
||||||
|
|||||||
@@ -29,7 +29,8 @@ use aether_data_contracts::repository::video_tasks::{
|
|||||||
UpsertVideoTask, VideoTaskLookupKey, VideoTaskStatus, VideoTaskWriteRepository,
|
UpsertVideoTask, VideoTaskLookupKey, VideoTaskStatus, VideoTaskWriteRepository,
|
||||||
};
|
};
|
||||||
use aether_scheduler_core::{
|
use aether_scheduler_core::{
|
||||||
build_minimal_candidate_selection, SchedulerAuthConstraints, SchedulerPriorityMode,
|
build_minimal_candidate_selection, BuildMinimalCandidateSelectionInput,
|
||||||
|
SchedulerAuthConstraints, SchedulerPriorityMode,
|
||||||
};
|
};
|
||||||
use async_trait::async_trait;
|
use async_trait::async_trait;
|
||||||
use serde_json::json;
|
use serde_json::json;
|
||||||
@@ -664,17 +665,17 @@ async fn data_state_reads_minimal_candidate_selection_with_auth_filters() {
|
|||||||
.map(|items| items.to_vec()),
|
.map(|items| items.to_vec()),
|
||||||
};
|
};
|
||||||
|
|
||||||
let selection = build_minimal_candidate_selection(
|
let selection = build_minimal_candidate_selection(BuildMinimalCandidateSelectionInput {
|
||||||
rows,
|
rows,
|
||||||
"openai:chat",
|
normalized_api_format: "openai:chat",
|
||||||
"gpt-4.1",
|
requested_model_name: "gpt-4.1",
|
||||||
"gpt-4.1",
|
resolved_global_model_name: "gpt-4.1",
|
||||||
false,
|
require_streaming: false,
|
||||||
None,
|
required_capabilities: None,
|
||||||
Some(&auth_constraints),
|
auth_constraints: Some(&auth_constraints),
|
||||||
Some(auth_snapshot.api_key_id.as_str()),
|
affinity_key: Some(auth_snapshot.api_key_id.as_str()),
|
||||||
SchedulerPriorityMode::Provider,
|
priority_mode: SchedulerPriorityMode::Provider,
|
||||||
)
|
})
|
||||||
.expect("selection should read");
|
.expect("selection should read");
|
||||||
|
|
||||||
assert_eq!(selection.len(), 2);
|
assert_eq!(selection.len(), 2);
|
||||||
|
|||||||
@@ -790,7 +790,6 @@ async fn execute_stream_from_frame_stream(
|
|||||||
let request_id_for_report = request_id.to_string();
|
let request_id_for_report = request_id.to_string();
|
||||||
let request_id_for_report_log = short_request_id(request_id);
|
let request_id_for_report_log = short_request_id(request_id);
|
||||||
let candidate_id_for_report = candidate_id.map(ToOwned::to_owned);
|
let candidate_id_for_report = candidate_id.map(ToOwned::to_owned);
|
||||||
let mut buffered_frames = buffered_frames;
|
|
||||||
tokio::spawn(async move {
|
tokio::spawn(async move {
|
||||||
let mut provider_buffered_body = provider_prefetched_body_for_report;
|
let mut provider_buffered_body = provider_prefetched_body_for_report;
|
||||||
let mut buffered_body = prefetched_body_for_report;
|
let mut buffered_body = prefetched_body_for_report;
|
||||||
|
|||||||
@@ -23,7 +23,10 @@ pub(crate) fn build_admin_global_model_response(
|
|||||||
"provider_count": global_model.provider_count,
|
"provider_count": global_model.provider_count,
|
||||||
"active_provider_count": global_model.active_provider_count,
|
"active_provider_count": global_model.active_provider_count,
|
||||||
"usage_count": global_model.usage_count,
|
"usage_count": global_model.usage_count,
|
||||||
"created_at": timestamp_or_now(global_model.created_at_unix_secs, now_unix_secs),
|
"created_at": timestamp_or_now(
|
||||||
|
global_model.created_at_unix_ms.map(|value| value / 1000),
|
||||||
|
now_unix_secs,
|
||||||
|
),
|
||||||
"updated_at": timestamp_or_now(global_model.updated_at_unix_secs, now_unix_secs),
|
"updated_at": timestamp_or_now(global_model.updated_at_unix_secs, now_unix_secs),
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -2,19 +2,14 @@ use aether_scheduler_core::SchedulerPriorityMode;
|
|||||||
|
|
||||||
use crate::{AppState, GatewayError};
|
use crate::{AppState, GatewayError};
|
||||||
|
|
||||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
|
||||||
pub(crate) enum SchedulerSchedulingMode {
|
pub(crate) enum SchedulerSchedulingMode {
|
||||||
FixedOrder,
|
FixedOrder,
|
||||||
|
#[default]
|
||||||
CacheAffinity,
|
CacheAffinity,
|
||||||
LoadBalance,
|
LoadBalance,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl Default for SchedulerSchedulingMode {
|
|
||||||
fn default() -> Self {
|
|
||||||
Self::CacheAffinity
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||||
pub(crate) struct SchedulerOrderingConfig {
|
pub(crate) struct SchedulerOrderingConfig {
|
||||||
pub(crate) priority_mode: SchedulerPriorityMode,
|
pub(crate) priority_mode: SchedulerPriorityMode,
|
||||||
|
|||||||
@@ -441,8 +441,7 @@ fn maybe_build_openai_cli_same_family_sync_body(
|
|||||||
|
|
||||||
if !is_openai_cli_family_api_format(&provider_api_format)
|
if !is_openai_cli_family_api_format(&provider_api_format)
|
||||||
|| !is_openai_cli_family_api_format(&client_api_format)
|
|| !is_openai_cli_family_api_format(&client_api_format)
|
||||||
|| provider_api_format != client_api_format
|
|| (provider_api_format == client_api_format && needs_conversion)
|
||||||
|| needs_conversion
|
|
||||||
{
|
{
|
||||||
return None;
|
return None;
|
||||||
}
|
}
|
||||||
@@ -488,8 +487,7 @@ fn maybe_build_openai_cli_same_family_stream_sync_body(
|
|||||||
|
|
||||||
if !is_openai_cli_family_api_format(&provider_api_format)
|
if !is_openai_cli_family_api_format(&provider_api_format)
|
||||||
|| !is_openai_cli_family_api_format(&client_api_format)
|
|| !is_openai_cli_family_api_format(&client_api_format)
|
||||||
|| provider_api_format != client_api_format
|
|| (provider_api_format == client_api_format && needs_conversion)
|
||||||
|| needs_conversion
|
|
||||||
{
|
{
|
||||||
return Ok(None);
|
return Ok(None);
|
||||||
}
|
}
|
||||||
@@ -1551,6 +1549,33 @@ mod tests {
|
|||||||
assert_eq!(body_json, provider_body_json);
|
assert_eq!(body_json, provider_body_json);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn allows_openai_cli_same_family_cross_format_sync_when_conversion_is_flagged() {
|
||||||
|
let report_context = json!({
|
||||||
|
"provider_api_format": "openai:compact",
|
||||||
|
"client_api_format": "openai:cli",
|
||||||
|
"needs_conversion": true,
|
||||||
|
});
|
||||||
|
let provider_body_json = json!({
|
||||||
|
"id": "resp_family_123",
|
||||||
|
"object": "response",
|
||||||
|
"status": "completed",
|
||||||
|
"output": []
|
||||||
|
});
|
||||||
|
|
||||||
|
let body_json = maybe_build_openai_cli_same_family_sync_body_from_normalized_payload(
|
||||||
|
"openai_cli_sync_finalize",
|
||||||
|
200,
|
||||||
|
Some(&report_context),
|
||||||
|
Some(&provider_body_json),
|
||||||
|
None,
|
||||||
|
)
|
||||||
|
.expect("openai-cli cross-family sync should succeed")
|
||||||
|
.expect("body should exist");
|
||||||
|
|
||||||
|
assert_eq!(body_json, provider_body_json);
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn rejects_openai_cli_same_family_error_body_json() {
|
fn rejects_openai_cli_same_family_error_body_json() {
|
||||||
let report_context = json!({
|
let report_context = json!({
|
||||||
@@ -1816,6 +1841,37 @@ mod tests {
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn standard_sync_finalize_product_handles_openai_cli_same_family_cross_format_body() {
|
||||||
|
let report_context = json!({
|
||||||
|
"provider_api_format": "openai:compact",
|
||||||
|
"client_api_format": "openai:cli",
|
||||||
|
"needs_conversion": true,
|
||||||
|
});
|
||||||
|
let provider_body_json = json!({
|
||||||
|
"id": "resp_family_123",
|
||||||
|
"object": "response",
|
||||||
|
"status": "completed",
|
||||||
|
"output": []
|
||||||
|
});
|
||||||
|
|
||||||
|
let product = maybe_build_standard_sync_finalize_product_from_normalized_payload(
|
||||||
|
"openai_cli_sync_finalize",
|
||||||
|
200,
|
||||||
|
Some(&report_context),
|
||||||
|
Some(&provider_body_json),
|
||||||
|
None,
|
||||||
|
)
|
||||||
|
.expect("dispatch should succeed");
|
||||||
|
|
||||||
|
assert_eq!(
|
||||||
|
product,
|
||||||
|
Some(StandardSyncFinalizeNormalizedProduct::SuccessBody(
|
||||||
|
provider_body_json
|
||||||
|
))
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn standard_sync_finalize_product_handles_openai_chat_cross_format() {
|
fn standard_sync_finalize_product_handles_openai_chat_cross_format() {
|
||||||
let report_context = json!({
|
let report_context = json!({
|
||||||
|
|||||||
@@ -257,11 +257,15 @@ mod tests {
|
|||||||
provider_name: Some("OpenAI".to_string()),
|
provider_name: Some("OpenAI".to_string()),
|
||||||
provider_website: None,
|
provider_website: None,
|
||||||
provider_type: Some("custom".to_string()),
|
provider_type: Some("custom".to_string()),
|
||||||
|
provider_priority: Some(0),
|
||||||
|
provider_keep_priority_on_conversion: Some(false),
|
||||||
endpoint_api_format: Some("openai:chat".to_string()),
|
endpoint_api_format: Some("openai:chat".to_string()),
|
||||||
endpoint_api_family: Some("openai".to_string()),
|
endpoint_api_family: Some("openai".to_string()),
|
||||||
endpoint_kind: Some("chat".to_string()),
|
endpoint_kind: Some("chat".to_string()),
|
||||||
provider_key_name: Some("prod".to_string()),
|
provider_key_name: Some("prod".to_string()),
|
||||||
provider_key_auth_type: Some("api_key".to_string()),
|
provider_key_auth_type: Some("api_key".to_string()),
|
||||||
|
provider_key_internal_priority: Some(10),
|
||||||
|
provider_key_global_priority_by_format: None,
|
||||||
provider_key_capabilities: None,
|
provider_key_capabilities: None,
|
||||||
provider_key_is_active: Some(true),
|
provider_key_is_active: Some(true),
|
||||||
}],
|
}],
|
||||||
|
|||||||
@@ -10,6 +10,16 @@ use super::{
|
|||||||
};
|
};
|
||||||
use crate::DataLayerError;
|
use crate::DataLayerError;
|
||||||
|
|
||||||
|
const MILLIS_PER_SECOND: u64 = 1000;
|
||||||
|
|
||||||
|
fn unix_secs_to_ms(unix_secs: u64) -> u64 {
|
||||||
|
unix_secs.saturating_mul(MILLIS_PER_SECOND)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn unix_ms_to_secs(unix_ms: u64) -> u64 {
|
||||||
|
unix_ms / MILLIS_PER_SECOND
|
||||||
|
}
|
||||||
|
|
||||||
#[derive(Debug, Default)]
|
#[derive(Debug, Default)]
|
||||||
pub struct InMemoryUsageReadRepository {
|
pub struct InMemoryUsageReadRepository {
|
||||||
by_request_id: RwLock<BTreeMap<String, StoredRequestUsageAudit>>,
|
by_request_id: RwLock<BTreeMap<String, StoredRequestUsageAudit>>,
|
||||||
@@ -80,12 +90,12 @@ impl UsageReadRepository for InMemoryUsageReadRepository {
|
|||||||
.values()
|
.values()
|
||||||
.filter(|item| {
|
.filter(|item| {
|
||||||
if let Some(created_from_unix_secs) = query.created_from_unix_secs {
|
if let Some(created_from_unix_secs) = query.created_from_unix_secs {
|
||||||
if item.created_at_unix_ms < created_from_unix_secs {
|
if item.created_at_unix_ms < unix_secs_to_ms(created_from_unix_secs) {
|
||||||
return false;
|
return false;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
if let Some(created_until_unix_secs) = query.created_until_unix_secs {
|
if let Some(created_until_unix_secs) = query.created_until_unix_secs {
|
||||||
if item.created_at_unix_ms >= created_until_unix_secs {
|
if item.created_at_unix_ms >= unix_secs_to_ms(created_until_unix_secs) {
|
||||||
return false;
|
return false;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -200,7 +210,7 @@ impl UsageReadRepository for InMemoryUsageReadRepository {
|
|||||||
entry
|
entry
|
||||||
.last_used_at_unix_secs
|
.last_used_at_unix_secs
|
||||||
.unwrap_or(0)
|
.unwrap_or(0)
|
||||||
.max(item.created_at_unix_secs),
|
.max(unix_ms_to_secs(item.created_at_unix_ms)),
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
Ok(summaries)
|
Ok(summaries)
|
||||||
@@ -558,4 +568,44 @@ mod tests {
|
|||||||
assert_eq!(summary.avg_response_time_ms, 180.0);
|
assert_eq!(summary.avg_response_time_ms, 180.0);
|
||||||
assert_eq!(summary.total_cost_usd, 0.75);
|
assert_eq!(summary.total_cost_usd, 0.75);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn list_usage_audits_applies_second_based_time_filters_to_millisecond_timestamps() {
|
||||||
|
let repository = InMemoryUsageReadRepository::seed(vec![
|
||||||
|
sample_usage("req-1", 1_000),
|
||||||
|
sample_usage("req-2", 2_000),
|
||||||
|
sample_usage("req-3", 3_000),
|
||||||
|
]);
|
||||||
|
|
||||||
|
let items = repository
|
||||||
|
.list_usage_audits(&crate::repository::usage::UsageAuditListQuery {
|
||||||
|
created_from_unix_secs: Some(2),
|
||||||
|
created_until_unix_secs: Some(3),
|
||||||
|
..Default::default()
|
||||||
|
})
|
||||||
|
.await
|
||||||
|
.expect("list should succeed");
|
||||||
|
|
||||||
|
assert_eq!(items.len(), 1);
|
||||||
|
assert_eq!(items[0].request_id, "req-2");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn summarizes_provider_api_key_last_used_at_in_seconds() {
|
||||||
|
let repository = InMemoryUsageReadRepository::seed(vec![
|
||||||
|
sample_usage("req-1", 1_999),
|
||||||
|
sample_usage("req-2", 2_500),
|
||||||
|
]);
|
||||||
|
|
||||||
|
let summary = repository
|
||||||
|
.summarize_usage_by_provider_api_key_ids(&["provider-key-1".to_string()])
|
||||||
|
.await
|
||||||
|
.expect("summary should succeed");
|
||||||
|
|
||||||
|
let usage = summary
|
||||||
|
.get("provider-key-1")
|
||||||
|
.expect("provider key summary should exist");
|
||||||
|
assert_eq!(usage.request_count, 2);
|
||||||
|
assert_eq!(usage.last_used_at_unix_secs, Some(2));
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -5,18 +5,13 @@ use aether_data_contracts::repository::candidates::StoredRequestCandidate;
|
|||||||
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
|
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
|
||||||
use aether_data_contracts::DataLayerError;
|
use aether_data_contracts::DataLayerError;
|
||||||
|
|
||||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, serde::Serialize, serde::Deserialize)]
|
||||||
pub enum SchedulerPriorityMode {
|
pub enum SchedulerPriorityMode {
|
||||||
|
#[default]
|
||||||
Provider,
|
Provider,
|
||||||
GlobalKey,
|
GlobalKey,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl Default for SchedulerPriorityMode {
|
|
||||||
fn default() -> Self {
|
|
||||||
Self::Provider
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[derive(Debug, Clone, PartialEq, serde::Serialize)]
|
#[derive(Debug, Clone, PartialEq, serde::Serialize)]
|
||||||
pub struct SchedulerMinimalCandidateSelectionCandidate {
|
pub struct SchedulerMinimalCandidateSelectionCandidate {
|
||||||
pub provider_id: String,
|
pub provider_id: String,
|
||||||
@@ -38,6 +33,18 @@ pub struct SchedulerMinimalCandidateSelectionCandidate {
|
|||||||
pub mapping_matched_model: Option<String>,
|
pub mapping_matched_model: Option<String>,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub struct BuildMinimalCandidateSelectionInput<'a> {
|
||||||
|
pub rows: Vec<StoredMinimalCandidateSelectionRow>,
|
||||||
|
pub normalized_api_format: &'a str,
|
||||||
|
pub requested_model_name: &'a str,
|
||||||
|
pub resolved_global_model_name: &'a str,
|
||||||
|
pub require_streaming: bool,
|
||||||
|
pub required_capabilities: Option<&'a serde_json::Value>,
|
||||||
|
pub auth_constraints: Option<&'a crate::SchedulerAuthConstraints>,
|
||||||
|
pub affinity_key: Option<&'a str>,
|
||||||
|
pub priority_mode: SchedulerPriorityMode,
|
||||||
|
}
|
||||||
|
|
||||||
pub fn candidate_supports_required_capability(
|
pub fn candidate_supports_required_capability(
|
||||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||||
required_capability: &str,
|
required_capability: &str,
|
||||||
@@ -118,16 +125,20 @@ pub fn auth_api_key_concurrency_limit_reached(
|
|||||||
}
|
}
|
||||||
|
|
||||||
pub fn build_minimal_candidate_selection(
|
pub fn build_minimal_candidate_selection(
|
||||||
rows: Vec<StoredMinimalCandidateSelectionRow>,
|
input: BuildMinimalCandidateSelectionInput<'_>,
|
||||||
normalized_api_format: &str,
|
|
||||||
requested_model_name: &str,
|
|
||||||
resolved_global_model_name: &str,
|
|
||||||
require_streaming: bool,
|
|
||||||
required_capabilities: Option<&serde_json::Value>,
|
|
||||||
auth_constraints: Option<&crate::SchedulerAuthConstraints>,
|
|
||||||
affinity_key: Option<&str>,
|
|
||||||
priority_mode: SchedulerPriorityMode,
|
|
||||||
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, DataLayerError> {
|
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, DataLayerError> {
|
||||||
|
let BuildMinimalCandidateSelectionInput {
|
||||||
|
rows,
|
||||||
|
normalized_api_format,
|
||||||
|
requested_model_name,
|
||||||
|
resolved_global_model_name,
|
||||||
|
require_streaming,
|
||||||
|
required_capabilities,
|
||||||
|
auth_constraints,
|
||||||
|
affinity_key,
|
||||||
|
priority_mode,
|
||||||
|
} = input;
|
||||||
|
|
||||||
if normalized_api_format.is_empty() {
|
if normalized_api_format.is_empty() {
|
||||||
return Ok(Vec::new());
|
return Ok(Vec::new());
|
||||||
}
|
}
|
||||||
@@ -497,8 +508,8 @@ mod tests {
|
|||||||
candidate_is_selectable_with_runtime_state, candidate_supports_required_capability,
|
candidate_is_selectable_with_runtime_state, candidate_supports_required_capability,
|
||||||
collect_global_model_names_for_required_capability,
|
collect_global_model_names_for_required_capability,
|
||||||
collect_selectable_candidates_from_keys, reorder_candidates_by_scheduler_health,
|
collect_selectable_candidates_from_keys, reorder_candidates_by_scheduler_health,
|
||||||
CandidateRuntimeSelectabilityInput, SchedulerMinimalCandidateSelectionCandidate,
|
BuildMinimalCandidateSelectionInput, CandidateRuntimeSelectabilityInput,
|
||||||
SchedulerPriorityMode,
|
SchedulerMinimalCandidateSelectionCandidate, SchedulerPriorityMode,
|
||||||
};
|
};
|
||||||
use crate::SchedulerAuthConstraints;
|
use crate::SchedulerAuthConstraints;
|
||||||
|
|
||||||
@@ -647,17 +658,17 @@ mod tests {
|
|||||||
allowed_api_formats: Some(vec!["OPENAI:CHAT".to_string()]),
|
allowed_api_formats: Some(vec!["OPENAI:CHAT".to_string()]),
|
||||||
allowed_models: Some(vec!["gpt-5".to_string()]),
|
allowed_models: Some(vec!["gpt-5".to_string()]),
|
||||||
};
|
};
|
||||||
let candidates = build_minimal_candidate_selection(
|
let candidates = build_minimal_candidate_selection(BuildMinimalCandidateSelectionInput {
|
||||||
vec![sample_row("1"), disallowed],
|
rows: vec![sample_row("1"), disallowed],
|
||||||
"openai:chat",
|
normalized_api_format: "openai:chat",
|
||||||
"gpt-5",
|
requested_model_name: "gpt-5",
|
||||||
"gpt-5",
|
resolved_global_model_name: "gpt-5",
|
||||||
false,
|
require_streaming: false,
|
||||||
None,
|
required_capabilities: None,
|
||||||
Some(&constraints),
|
auth_constraints: Some(&constraints),
|
||||||
None,
|
affinity_key: None,
|
||||||
SchedulerPriorityMode::Provider,
|
priority_mode: SchedulerPriorityMode::Provider,
|
||||||
)
|
})
|
||||||
.expect("candidate selection should build");
|
.expect("candidate selection should build");
|
||||||
|
|
||||||
assert_eq!(candidates.len(), 1);
|
assert_eq!(candidates.len(), 1);
|
||||||
@@ -699,17 +710,17 @@ mod tests {
|
|||||||
matching_capability.provider_priority = 10;
|
matching_capability.provider_priority = 10;
|
||||||
|
|
||||||
let required_capabilities = serde_json::json!({"cache_1h": true});
|
let required_capabilities = serde_json::json!({"cache_1h": true});
|
||||||
let candidates = build_minimal_candidate_selection(
|
let candidates = build_minimal_candidate_selection(BuildMinimalCandidateSelectionInput {
|
||||||
vec![missing_capability, matching_capability],
|
rows: vec![missing_capability, matching_capability],
|
||||||
"openai:chat",
|
normalized_api_format: "openai:chat",
|
||||||
"gpt-5",
|
requested_model_name: "gpt-5",
|
||||||
"gpt-5",
|
resolved_global_model_name: "gpt-5",
|
||||||
false,
|
require_streaming: false,
|
||||||
Some(&required_capabilities),
|
required_capabilities: Some(&required_capabilities),
|
||||||
None,
|
auth_constraints: None,
|
||||||
None,
|
affinity_key: None,
|
||||||
SchedulerPriorityMode::Provider,
|
priority_mode: SchedulerPriorityMode::Provider,
|
||||||
)
|
})
|
||||||
.expect("candidate selection should build");
|
.expect("candidate selection should build");
|
||||||
|
|
||||||
assert_eq!(candidates.len(), 2);
|
assert_eq!(candidates.len(), 2);
|
||||||
|
|||||||
@@ -19,8 +19,9 @@ pub use candidate::{
|
|||||||
candidate_is_selectable_with_runtime_state, candidate_supports_required_capability,
|
candidate_is_selectable_with_runtime_state, candidate_supports_required_capability,
|
||||||
collect_global_model_names_for_required_capability, collect_selectable_candidates_from_keys,
|
collect_global_model_names_for_required_capability, collect_selectable_candidates_from_keys,
|
||||||
compare_candidates_by_priority_mode, reorder_candidates_by_scheduler_health,
|
compare_candidates_by_priority_mode, reorder_candidates_by_scheduler_health,
|
||||||
requested_capability_priority_for_candidate, CandidateRuntimeSelectabilityInput,
|
requested_capability_priority_for_candidate, BuildMinimalCandidateSelectionInput,
|
||||||
SchedulerMinimalCandidateSelectionCandidate, SchedulerPriorityMode,
|
CandidateRuntimeSelectabilityInput, SchedulerMinimalCandidateSelectionCandidate,
|
||||||
|
SchedulerPriorityMode,
|
||||||
};
|
};
|
||||||
pub use health::{
|
pub use health::{
|
||||||
aggregate_provider_key_health_score, count_recent_active_requests_for_api_key,
|
aggregate_provider_key_health_score, count_recent_active_requests_for_api_key,
|
||||||
|
|||||||
Reference in New Issue
Block a user