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

feat: add model directive management

Browse Source
This commit is contained in:
fawney19
2026-05-03 14:48:25 +08:00
parent fe27fb17fb
commit c4ea042eb4
53 changed files with 2655 additions and 182 deletions
Download Patch File Download Diff File Expand all files Collapse all files

62
crates/aether-admin/src/system.rs
View File

@@ -983,7 +983,7 @@ pub fn build_admin_module_validation_result(
)
}
}
"management_tokens" | "proxy_nodes" => (true, None),
"management_tokens" | "model_directives" | "proxy_nodes" => (true, None),
_ => (true, None),
}
}
@@ -993,7 +993,7 @@ pub fn build_admin_module_health(
gemini_files_has_capable_key: bool,
) -> &'static str {
match module_name {
"management_tokens" | "proxy_nodes" => "healthy",
"management_tokens" | "model_directives" | "proxy_nodes" => "healthy",
"gemini_files" => {
if gemini_files_has_capable_key {
"healthy"
@@ -1225,6 +1225,64 @@ pub fn admin_system_config_default_value(key: &str) -> Option<serde_json::Value>
"email_suffix_mode" => Some(json!("none")),
"email_suffix_list" => Some(json!([])),
"enable_format_conversion" => Some(json!(false)),
"enable_model_directives" => Some(json!(false)),
"model_directives" => Some(json!({
"reasoning_effort": {
"enabled": true,
"api_formats": {
"openai:chat": {
"enabled": true,
"mappings": {
"low": { "reasoning_effort": "low" },
"medium": { "reasoning_effort": "medium" },
"high": { "reasoning_effort": "high" },
"xhigh": { "reasoning_effort": "xhigh" },
"max": { "reasoning_effort": "xhigh" }
}
},
"openai:responses": {
"enabled": true,
"mappings": {
"low": { "reasoning": { "effort": "low" } },
"medium": { "reasoning": { "effort": "medium" } },
"high": { "reasoning": { "effort": "high" } },
"xhigh": { "reasoning": { "effort": "xhigh" } },
"max": { "reasoning": { "effort": "xhigh" } }
}
},
"openai:responses:compact": {
"enabled": true,
"mappings": {
"low": { "reasoning": { "effort": "low" } },
"medium": { "reasoning": { "effort": "medium" } },
"high": { "reasoning": { "effort": "high" } },
"xhigh": { "reasoning": { "effort": "xhigh" } },
"max": { "reasoning": { "effort": "xhigh" } }
}
},
"claude:messages": {
"enabled": true,
"mappings": {
"low": { "thinking": { "type": "enabled", "budget_tokens": 1024 } },
"medium": { "thinking": { "type": "enabled", "budget_tokens": 4096 } },
"high": { "thinking": { "type": "enabled", "budget_tokens": 8192 } },
"xhigh": { "thinking": { "type": "enabled", "budget_tokens": 16384 } },
"max": { "thinking": { "type": "enabled", "budget_tokens": 32768 } }
}
},
"gemini:generate_content": {
"enabled": true,
"mappings": {
"low": { "generationConfig": { "thinkingConfig": { "thinkingBudget": 1024 } } },
"medium": { "generationConfig": { "thinkingConfig": { "thinkingBudget": 4096 } } },
"high": { "generationConfig": { "thinkingConfig": { "thinkingBudget": 8192 } } },
"xhigh": { "generationConfig": { "thinkingConfig": { "thinkingBudget": 16384 } } },
"max": { "generationConfig": { "thinkingConfig": { "thinkingBudget": -1 } } }
}
}
}
}
})),
"keep_priority_on_conversion" => Some(json!(false)),
"audit_log_retention_days" => Some(json!(30)),
"enable_db_maintenance" => Some(json!(true)),

22
crates/aether-ai-formats/src/api.rs
View File

@@ -61,7 +61,17 @@ pub use crate::provider_compat::surfaces::{
pub use crate::request::common::{
force_upstream_streaming_for_provider, parse_direct_request_body,
};
pub use crate::request::matrix::build_standard_request_body_from_canonical;
pub use crate::request::matrix::{
build_standard_request_body_from_canonical,
build_standard_request_body_from_canonical_with_model_directives,
};
pub use crate::request::model_directives::{
apply_model_directive_mapping_patch, apply_model_directive_overrides_from_model,
apply_model_directive_overrides_from_request, claude_model_uses_adaptive_effort,
extract_gemini_model_from_path, gemini_model_uses_thinking_level, model_directive_base_model,
normalize_model_directive_model, parse_model_directive, ModelDirective, ModelOverride,
ReasoningEffort,
};
pub use crate::request::openai::{
copy_request_number_field, copy_request_number_field_as,
map_openai_reasoning_effort_to_claude_output, map_openai_reasoning_effort_to_gemini_budget,
@@ -98,8 +108,14 @@ pub use crate::request::specialized::{
pub use crate::request::standard::{
apply_codex_openai_responses_special_body_edits, apply_codex_openai_responses_special_headers,
apply_openai_responses_compact_special_body_edits, build_cross_format_openai_chat_request_body,
build_cross_format_openai_responses_request_body, build_local_openai_chat_request_body,
build_local_openai_responses_request_body, build_standard_request_body,
build_cross_format_openai_chat_request_body_with_model_directives,
build_cross_format_openai_responses_request_body,
build_cross_format_openai_responses_request_body_with_model_directives,
build_local_openai_chat_request_body,
build_local_openai_chat_request_body_with_model_directives,
build_local_openai_responses_request_body,
build_local_openai_responses_request_body_with_model_directives, build_standard_request_body,
build_standard_request_body_with_model_directives,
claude::{
resolve_stream_spec as resolve_claude_stream_spec,
resolve_sync_spec as resolve_claude_sync_spec,

7
crates/aether-ai-formats/src/lib.rs
View File

@@ -36,3 +36,10 @@ pub use protocol::matrix::{
SyncChatResponseConversionKind, SyncCliResponseConversionKind,
};
pub use protocol::registry::{build_stream_transcoder, convert_request, convert_response};
pub use request::model_directives::{
apply_model_directive_mapping_patch, apply_model_directive_overrides_from_model,
apply_model_directive_overrides_from_request, claude_model_uses_adaptive_effort,
extract_gemini_model_from_path, gemini_model_uses_thinking_level, model_directive_base_model,
normalize_model_directive_model, parse_model_directive, ModelDirective, ModelOverride,
ReasoningEffort,
};

3
crates/aether-ai-formats/src/protocol/canonical.rs
View File

@@ -2503,7 +2503,8 @@ pub(crate) fn claude_output_effort_to_openai_reasoning_effort(value: &str) -> Op
"low" => Some("low"),
"medium" => Some("medium"),
"high" => Some("high"),
"max" | "xhigh" => Some("xhigh"),
"xhigh" => Some("xhigh"),
"max" => Some("max"),
_ => None,
}
}

21
crates/aether-ai-formats/src/protocol/formats/claude_messages/request.rs
View File

@@ -12,9 +12,12 @@ use crate::{
CanonicalRequest,
},
protocol::context::FormatContext,
request::openai::{
map_openai_reasoning_effort_to_claude_output,
map_openai_reasoning_effort_to_thinking_budget,
request::{
model_directives::claude_model_uses_adaptive_effort,
openai::{
map_openai_reasoning_effort_to_claude_output,
map_openai_reasoning_effort_to_thinking_budget,
},
},
};
@@ -155,14 +158,18 @@ pub fn to_raw(
let budget_tokens = thinking
.budget_tokens
.or_else(|| openai_effort.and_then(map_openai_reasoning_effort_to_thinking_budget));
let uses_adaptive = claude_model_uses_adaptive_effort(mapped_model)
|| claude_model_uses_adaptive_effort(canonical.model.as_str());
if thinking.enabled || budget_tokens.is_some() {
output.insert(
"thinking".to_string(),
let thinking_config = if uses_adaptive {
json!({"type": "adaptive"})
} else {
json!({
"type": "enabled",
"budget_tokens": budget_tokens.unwrap_or(1024),
}),
);
})
};
output.insert("thinking".to_string(), thinking_config);
}
if let Some(output_effort) =
openai_effort.and_then(map_openai_reasoning_effort_to_claude_output)

54
crates/aether-ai-formats/src/protocol/formats/gemini_generate_content/request.rs
View File

@@ -15,7 +15,13 @@ use crate::{
OPENAI_RESPONSES_LEGACY_EXTENSION_NAMESPACE,
},
protocol::context::FormatContext,
request::openai::map_openai_reasoning_effort_to_gemini_budget,
request::{
model_directives::{gemini_model_uses_thinking_level, ReasoningEffort},
openai::{
map_openai_reasoning_effort_to_gemini_budget,
map_thinking_budget_to_openai_reasoning_effort,
},
},
};
pub fn from(body: &Value, ctx: &FormatContext) -> Option<CanonicalRequest> {
@@ -203,7 +209,8 @@ fn canonical_to_gemini_request_body(
if let Some(system_instruction) = canonical_system_instruction(canonical) {
output.insert("systemInstruction".to_string(), system_instruction);
}
if let Some(generation_config) = canonical_generation_config_to_gemini(canonical) {
if let Some(generation_config) = canonical_generation_config_to_gemini(canonical, mapped_model)
{
output.insert("generationConfig".to_string(), generation_config);
}
if let Some(tools) = canonical_tools_to_gemini(canonical) {
@@ -370,7 +377,10 @@ fn canonical_media_to_gemini_part(
})
}
fn canonical_generation_config_to_gemini(canonical: &CanonicalRequest) -> Option<Value> {
fn canonical_generation_config_to_gemini(
canonical: &CanonicalRequest,
mapped_model: &str,
) -> Option<Value> {
let mut generation_config = Map::new();
if let Some(value) = canonical.generation.max_tokens {
generation_config.insert("maxOutputTokens".to_string(), Value::from(value));
@@ -406,14 +416,8 @@ fn canonical_generation_config_to_gemini(canonical: &CanonicalRequest) -> Option
.and_then(|value| value.get("thinking_config"))
.cloned()
.or_else(|| {
let budget = thinking.budget_tokens.or_else(|| {
canonical_openai_reasoning_effort(thinking)
.and_then(map_openai_reasoning_effort_to_gemini_budget)
})?;
Some(json!({
"includeThoughts": true,
"thinkingBudget": budget,
}))
let effort = canonical_openai_reasoning_effort(thinking);
gemini_thinking_config_from_reasoning(mapped_model, effort, thinking.budget_tokens)
})
}) {
generation_config.insert("thinkingConfig".to_string(), thinking_config);
@@ -421,6 +425,34 @@ fn canonical_generation_config_to_gemini(canonical: &CanonicalRequest) -> Option
(!generation_config.is_empty()).then_some(Value::Object(generation_config))
}
fn gemini_thinking_config_from_reasoning(
mapped_model: &str,
effort: Option<&str>,
budget_tokens: Option<u64>,
) -> Option<Value> {
if gemini_model_uses_thinking_level(mapped_model) {
let level = effort
.and_then(ReasoningEffort::parse)
.or_else(|| {
budget_tokens
.map(map_thinking_budget_to_openai_reasoning_effort)
.and_then(ReasoningEffort::parse)
})
.map(ReasoningEffort::as_gemini_level_value)?;
return Some(json!({
"includeThoughts": true,
"thinkingLevel": level,
}));
}
let budget =
budget_tokens.or_else(|| effort.and_then(map_openai_reasoning_effort_to_gemini_budget))?;
Some(json!({
"includeThoughts": true,
"thinkingBudget": budget,
}))
}
fn apply_response_format_to_gemini_generation_config(
generation_config: &mut Map<String, Value>,
response_format: &CanonicalResponseFormat,

12
crates/aether-ai-formats/src/protocol/formats/openai_responses/request.rs
View File

@@ -397,7 +397,7 @@ fn reasoning_config_to_responses(thinking: &CanonicalThinkingConfig) -> Option<V
.and_then(Value::as_str)
.map(|effort| {
json!({
"effort": if effort == "xhigh" { "high" } else { effort },
"effort": openai_responses_reasoning_effort(effort),
})
})
})
@@ -410,6 +410,16 @@ fn reasoning_config_to_responses(thinking: &CanonicalThinkingConfig) -> Option<V
})
}
fn openai_responses_reasoning_effort(effort: &str) -> &str {
match effort.trim().to_ascii_lowercase().as_str() {
"xhigh" | "max" => "xhigh",
"low" => "low",
"medium" => "medium",
"high" => "high",
_ => effort,
}
}
fn canonical_text_config_to_responses(canonical: &CanonicalRequest) -> Option<Value> {
let mut text = Map::new();
if let Some(response_format) = &canonical.response_format {

5
crates/aether-ai-formats/src/request/matrix.rs
View File

@@ -1 +1,4 @@
pub use crate::request::standard::matrix::build_standard_request_body_from_canonical;
pub use crate::request::standard::matrix::{
build_standard_request_body_from_canonical,
build_standard_request_body_from_canonical_with_model_directives,
};

1
crates/aether-ai-formats/src/request/mod.rs
View File

@@ -1,5 +1,6 @@
pub mod common;
pub mod matrix;
pub mod model_directives;
pub mod openai;
pub mod passthrough;
pub mod route;

434
crates/aether-ai-formats/src/request/model_directives.rs Normal file
View File

@@ -0,0 +1,434 @@
use serde_json::{json, Value};
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ModelDirective {
pub base_model: String,
pub overrides: Vec<ModelOverride>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ModelOverride {
ReasoningEffort(ReasoningEffort),
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ReasoningEffort {
Low,
Medium,
High,
XHigh,
Max,
}
impl ReasoningEffort {
pub fn parse(value: &str) -> Option<Self> {
match value.trim().to_ascii_lowercase().as_str() {
"low" => Some(Self::Low),
"medium" => Some(Self::Medium),
"high" => Some(Self::High),
"xhigh" => Some(Self::XHigh),
"max" => Some(Self::Max),
_ => None,
}
}
pub fn as_openai_chat_value(self) -> &'static str {
match self {
Self::Low => "low",
Self::Medium => "medium",
Self::High => "high",
Self::XHigh => "xhigh",
Self::Max => "xhigh",
}
}
pub fn as_openai_responses_value(self) -> &'static str {
match self {
Self::Low => "low",
Self::Medium => "medium",
Self::High => "high",
Self::XHigh | Self::Max => "xhigh",
}
}
pub fn as_claude_output_value(self) -> &'static str {
match self {
Self::Low => "low",
Self::Medium => "medium",
Self::High => "high",
Self::XHigh => "xhigh",
Self::Max => "max",
}
}
pub fn as_gemini_level_value(self) -> &'static str {
match self {
Self::Low => "low",
Self::Medium => "medium",
Self::High | Self::XHigh | Self::Max => "high",
}
}
pub fn thinking_budget_tokens(self) -> u64 {
match self {
Self::Low => 1280,
Self::Medium => 2048,
Self::High => 4096,
Self::XHigh | Self::Max => 8192,
}
}
}
pub fn parse_model_directive(model: &str) -> Option<ModelDirective> {
let model = model.trim();
let (base_model, suffix) = model.rsplit_once('-')?;
let base_model = base_model.trim();
if base_model.is_empty() {
return None;
}
let reasoning_effort = ReasoningEffort::parse(suffix)?;
Some(ModelDirective {
base_model: base_model.to_string(),
overrides: vec![ModelOverride::ReasoningEffort(reasoning_effort)],
})
}
pub fn model_directive_base_model(model: &str) -> Option<String> {
parse_model_directive(model).map(|directive| directive.base_model)
}
pub fn normalize_model_directive_model(model: &str) -> String {
parse_model_directive(model)
.map(|directive| directive.base_model)
.unwrap_or_else(|| model.trim().to_string())
}
pub fn apply_model_directive_overrides_from_request(
provider_request_body: &mut Value,
provider_api_format: &str,
provider_model: &str,
request_body: &Value,
request_path: Option<&str>,
) -> Option<ModelDirective> {
let source_model = request_body
.get("model")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
.or_else(|| request_path.and_then(extract_gemini_model_from_path))?;
apply_model_directive_overrides_from_model(
provider_request_body,
provider_api_format,
provider_model,
&source_model,
)
}
pub fn apply_model_directive_overrides_from_model(
provider_request_body: &mut Value,
provider_api_format: &str,
provider_model: &str,
source_model: &str,
) -> Option<ModelDirective> {
let directive = parse_model_directive(source_model)?;
for override_item in &directive.overrides {
match override_item {
ModelOverride::ReasoningEffort(effort) => {
apply_reasoning_effort_override(
provider_request_body,
provider_api_format,
provider_model,
*effort,
)?;
}
}
}
Some(directive)
}
pub fn apply_model_directive_mapping_patch(
provider_request_body: &mut Value,
patch: &Value,
) -> Option<()> {
deep_merge_json(provider_request_body, patch);
Some(())
}
fn deep_merge_json(target: &mut Value, patch: &Value) {
match (target, patch) {
(Value::Object(target_object), Value::Object(patch_object)) => {
for (key, patch_value) in patch_object {
match target_object.get_mut(key) {
Some(target_value) => deep_merge_json(target_value, patch_value),
None => {
target_object.insert(key.clone(), patch_value.clone());
}
}
}
}
(target, patch) => {
*target = patch.clone();
}
}
}
fn apply_reasoning_effort_override(
provider_request_body: &mut Value,
provider_api_format: &str,
provider_model: &str,
effort: ReasoningEffort,
) -> Option<()> {
match crate::normalize_api_format_alias(provider_api_format).as_str() {
"openai:chat" => set_object_string(
provider_request_body,
"reasoning_effort",
effort.as_openai_chat_value(),
),
"openai:responses" | "openai:responses:compact" => {
set_openai_responses_reasoning_effort(provider_request_body, effort)
}
"claude:messages" => {
set_claude_reasoning_effort(provider_request_body, effort, provider_model)
}
"gemini:generate_content" => {
set_gemini_reasoning_effort(provider_request_body, effort, provider_model)
}
_ => None,
}
}
fn set_object_string(body: &mut Value, key: &str, value: &str) -> Option<()> {
body.as_object_mut()?
.insert(key.to_string(), Value::String(value.to_string()));
Some(())
}
fn set_openai_responses_reasoning_effort(body: &mut Value, effort: ReasoningEffort) -> Option<()> {
let body_object = body.as_object_mut()?;
let reasoning = body_object
.entry("reasoning".to_string())
.or_insert_with(|| json!({}));
if !reasoning.is_object() {
*reasoning = json!({});
}
reasoning.as_object_mut()?.insert(
"effort".to_string(),
Value::String(effort.as_openai_responses_value().to_string()),
);
Some(())
}
fn set_claude_reasoning_effort(
body: &mut Value,
effort: ReasoningEffort,
provider_model: &str,
) -> Option<()> {
let body_object = body.as_object_mut()?;
let output_config = body_object
.entry("output_config".to_string())
.or_insert_with(|| json!({}));
if !output_config.is_object() {
*output_config = json!({});
}
output_config.as_object_mut()?.insert(
"effort".to_string(),
Value::String(effort.as_claude_output_value().to_string()),
);
let thinking = body_object
.entry("thinking".to_string())
.or_insert_with(|| json!({}));
if !thinking.is_object() {
*thinking = json!({});
}
let thinking = thinking.as_object_mut()?;
if claude_model_uses_adaptive_effort(provider_model) {
thinking.insert("type".to_string(), Value::String("adaptive".to_string()));
thinking.remove("budget_tokens");
} else {
thinking.insert("type".to_string(), Value::String("enabled".to_string()));
thinking.insert(
"budget_tokens".to_string(),
Value::from(effort.thinking_budget_tokens()),
);
}
Some(())
}
fn set_gemini_reasoning_effort(
body: &mut Value,
effort: ReasoningEffort,
provider_model: &str,
) -> Option<()> {
let body_object = body.as_object_mut()?;
let generation_key = if body_object.contains_key("generation_config")
&& !body_object.contains_key("generationConfig")
{
"generation_config"
} else {
"generationConfig"
};
let generation_config = body_object
.entry(generation_key.to_string())
.or_insert_with(|| json!({}));
if !generation_config.is_object() {
*generation_config = json!({});
}
let generation_config = generation_config.as_object_mut()?;
let thinking_key = if generation_config.contains_key("thinking_config")
&& !generation_config.contains_key("thinkingConfig")
{
"thinking_config"
} else {
"thinkingConfig"
};
generation_config.insert(
thinking_key.to_string(),
gemini_reasoning_effort_config(effort, provider_model, thinking_key),
);
Some(())
}
fn gemini_reasoning_effort_config(
effort: ReasoningEffort,
provider_model: &str,
thinking_key: &str,
) -> Value {
if gemini_model_uses_thinking_level(provider_model) {
if thinking_key == "thinking_config" {
return json!({
"include_thoughts": true,
"thinking_level": effort.as_gemini_level_value(),
});
}
return json!({
"includeThoughts": true,
"thinkingLevel": effort.as_gemini_level_value(),
});
}
if thinking_key == "thinking_config" {
return json!({
"include_thoughts": true,
"thinking_budget": effort.thinking_budget_tokens(),
});
}
json!({
"includeThoughts": true,
"thinkingBudget": effort.thinking_budget_tokens(),
})
}
pub fn claude_model_uses_adaptive_effort(model: &str) -> bool {
let model = model.trim().to_ascii_lowercase().replace(['.', '_'], "-");
model.contains("mythos")
|| model.contains("opus-4-7")
|| model.contains("opus-4-6")
|| model.contains("sonnet-4-6")
}
pub fn gemini_model_uses_thinking_level(model: &str) -> bool {
model
.trim()
.to_ascii_lowercase()
.split('/')
.any(|part| part.starts_with("gemini-3"))
}
pub fn extract_gemini_model_from_path(path: &str) -> Option<String> {
let marker = "/models/";
let start = path.find(marker)? + marker.len();
let tail = &path[start..];
let end = tail.find(':').unwrap_or(tail.len());
let model = tail[..end].trim();
(!model.is_empty()).then(|| model.to_string())
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::{
apply_model_directive_overrides_from_model, parse_model_directive, ModelDirective,
ModelOverride, ReasoningEffort,
};
#[test]
fn parses_supported_reasoning_effort_suffixes() {
assert_eq!(
parse_model_directive("gpt-5.4-xhigh"),
Some(ModelDirective {
base_model: "gpt-5.4".to_string(),
overrides: vec![ModelOverride::ReasoningEffort(ReasoningEffort::XHigh)],
})
);
assert_eq!(
parse_model_directive("gpt-5.4-MAX"),
Some(ModelDirective {
base_model: "gpt-5.4".to_string(),
overrides: vec![ModelOverride::ReasoningEffort(ReasoningEffort::Max)],
})
);
}
#[test]
fn ignores_unknown_or_incomplete_suffixes() {
assert_eq!(parse_model_directive("gpt-5.4-ultra"), None);
assert_eq!(parse_model_directive("gpt-5.4"), None);
assert_eq!(parse_model_directive("-high"), None);
assert_eq!(parse_model_directive("gpt-5.4-high-json"), None);
}
#[test]
fn applies_reasoning_effort_to_provider_body_shapes() {
let mut openai_chat = json!({"model": "gpt-5-upstream", "reasoning_effort": "low"});
apply_model_directive_overrides_from_model(
&mut openai_chat,
"openai:chat",
"gpt-5-upstream",
"gpt-5.4-xhigh",
)
.expect("directive should apply");
assert_eq!(openai_chat["reasoning_effort"], "xhigh");
let mut responses = json!({
"model": "gpt-5-upstream",
"reasoning": {"effort": "low", "summary": "auto"}
});
apply_model_directive_overrides_from_model(
&mut responses,
"openai:responses",
"gpt-5-upstream",
"gpt-5.4-max",
)
.expect("directive should apply");
assert_eq!(responses["reasoning"]["effort"], "xhigh");
assert_eq!(responses["reasoning"]["summary"], "auto");
let mut claude = json!({"model": "claude-sonnet-4-5"});
apply_model_directive_overrides_from_model(
&mut claude,
"claude:messages",
"claude-sonnet-4-5",
"gpt-5.4-high",
)
.expect("directive should apply");
assert_eq!(claude["thinking"]["budget_tokens"], 4096);
let mut gemini = json!({});
apply_model_directive_overrides_from_model(
&mut gemini,
"gemini:generate_content",
"gemini-2.5-pro",
"gpt-5.4-medium",
)
.expect("directive should apply");
assert_eq!(
gemini["generationConfig"]["thinkingConfig"]["thinkingBudget"],
2048
);
}
}

18
crates/aether-ai-formats/src/request/openai.rs
View File

@@ -1,5 +1,7 @@
use serde_json::{Map, Value};
use super::model_directives::ReasoningEffort;
pub fn parse_openai_stop_sequences(stop: Option<&Value>) -> Option<Vec<Value>> {
match stop {
Some(Value::String(value)) if !value.trim().is_empty() => {
@@ -55,23 +57,11 @@ pub fn copy_request_number_field_as(
}
pub fn map_openai_reasoning_effort_to_claude_output(value: &str) -> Option<&'static str> {
match value.trim().to_ascii_lowercase().as_str() {
"low" => Some("low"),
"medium" => Some("medium"),
"high" => Some("high"),
"xhigh" => Some("max"),
_ => None,
}
ReasoningEffort::parse(value).map(ReasoningEffort::as_claude_output_value)
}
pub fn map_openai_reasoning_effort_to_thinking_budget(value: &str) -> Option<u64> {
match value.trim().to_ascii_lowercase().as_str() {
"low" => Some(1280),
"medium" => Some(2048),
"high" => Some(4096),
"xhigh" => Some(8192),
_ => None,
}
ReasoningEffort::parse(value).map(ReasoningEffort::thinking_budget_tokens)
}
pub fn map_openai_reasoning_effort_to_gemini_budget(value: &str) -> Option<u64> {

178
crates/aether-ai-formats/src/request/standard/matrix.rs
View File

@@ -11,10 +11,12 @@ use aether_ai_formats::protocol::registry::{convert_request, FormatContext};
use aether_ai_formats::provider_compat::proxy::rules::apply_local_body_rules;
use serde_json::Value;
use crate::request::model_directives::apply_model_directive_overrides_from_request;
use super::{
apply_openai_responses_compact_special_body_edits,
codex::apply_codex_openai_responses_special_body_edits,
normalize::build_local_openai_chat_request_body,
normalize::build_local_openai_chat_request_body_with_model_directives,
};
#[allow(clippy::too_many_arguments)]
@@ -28,6 +30,33 @@ pub fn build_standard_request_body(
upstream_is_stream: bool,
body_rules: Option<&Value>,
user_api_key_id: Option<&str>,
) -> Option<Value> {
build_standard_request_body_with_model_directives(
body_json,
client_api_format,
mapped_model,
provider_type,
provider_api_format,
request_path,
upstream_is_stream,
body_rules,
user_api_key_id,
false,
)
}
#[allow(clippy::too_many_arguments)]
pub fn build_standard_request_body_with_model_directives(
body_json: &Value,
client_api_format: &str,
mapped_model: &str,
provider_type: &str,
provider_api_format: &str,
request_path: &str,
upstream_is_stream: bool,
body_rules: Option<&Value>,
user_api_key_id: Option<&str>,
enable_model_directives: bool,
) -> Option<Value> {
let format_context = FormatContext::default()
.with_mapped_model(mapped_model)
@@ -41,6 +70,16 @@ pub fn build_standard_request_body(
)
.ok()?;
if enable_model_directives {
apply_model_directive_overrides_from_request(
&mut provider_request_body,
provider_api_format,
mapped_model,
body_json,
Some(request_path),
);
}
if !apply_local_body_rules(&mut provider_request_body, body_rules, Some(body_json)) {
return None;
}
@@ -64,36 +103,64 @@ pub fn build_standard_request_body_from_canonical(
provider_api_format: &str,
upstream_is_stream: bool,
) -> Option<Value> {
match aether_ai_formats::normalize_api_format_alias(provider_api_format).as_str() {
"openai:chat" => build_local_openai_chat_request_body(
canonical_request,
build_standard_request_body_from_canonical_with_model_directives(
canonical_request,
mapped_model,
provider_api_format,
upstream_is_stream,
false,
)
}
pub fn build_standard_request_body_from_canonical_with_model_directives(
canonical_request: &Value,
mapped_model: &str,
provider_api_format: &str,
upstream_is_stream: bool,
enable_model_directives: bool,
) -> Option<Value> {
let mut provider_request_body =
match aether_ai_formats::normalize_api_format_alias(provider_api_format).as_str() {
"openai:chat" => build_local_openai_chat_request_body_with_model_directives(
canonical_request,
mapped_model,
upstream_is_stream,
enable_model_directives,
),
"openai:responses" => convert_openai_chat_request_to_openai_responses_request(
canonical_request,
mapped_model,
upstream_is_stream,
false,
),
"openai:responses:compact" => convert_openai_chat_request_to_openai_responses_request(
canonical_request,
mapped_model,
false,
true,
),
"claude:messages" => convert_openai_chat_request_to_claude_request(
canonical_request,
mapped_model,
upstream_is_stream,
),
"gemini:generate_content" => convert_openai_chat_request_to_gemini_request(
canonical_request,
mapped_model,
upstream_is_stream,
),
_ => None,
}?;
if enable_model_directives {
apply_model_directive_overrides_from_request(
&mut provider_request_body,
provider_api_format,
mapped_model,
upstream_is_stream,
),
"openai:responses" => convert_openai_chat_request_to_openai_responses_request(
canonical_request,
mapped_model,
upstream_is_stream,
false,
),
"openai:responses:compact" => convert_openai_chat_request_to_openai_responses_request(
canonical_request,
mapped_model,
false,
true,
),
"claude:messages" => convert_openai_chat_request_to_claude_request(
canonical_request,
mapped_model,
upstream_is_stream,
),
"gemini:generate_content" => convert_openai_chat_request_to_gemini_request(
canonical_request,
mapped_model,
upstream_is_stream,
),
_ => None,
None,
);
}
Some(provider_request_body)
}
pub fn normalize_standard_request_to_openai_chat_request(
@@ -133,6 +200,7 @@ fn normalize_standard_request_to_openai_chat_request_cow<'a>(
mod tests {
use super::{
build_standard_request_body, build_standard_request_body_from_canonical,
build_standard_request_body_with_model_directives,
normalize_standard_request_to_openai_chat_request,
};
use serde_json::{json, Value};
@@ -425,6 +493,60 @@ mod tests {
}
}
#[test]
fn standard_request_body_applies_reasoning_effort_suffix_to_claude_target() {
let request = json!({
"model": "gpt-5.4-max",
"messages": [{"role": "user", "content": "Need high effort"}],
"reasoning_effort": "low"
});
let converted = build_standard_request_body_with_model_directives(
&request,
"openai:chat",
"claude-sonnet-4-5",
"anthropic",
"claude:messages",
"/v1/chat/completions",
false,
None,
None,
true,
)
.expect("openai chat should convert to claude chat");
assert_eq!(converted["model"], "claude-sonnet-4-5");
assert_eq!(converted["output_config"]["effort"], "max");
assert_eq!(converted["thinking"]["budget_tokens"], 8192);
}
#[test]
fn standard_request_body_applies_reasoning_effort_suffix_from_gemini_path() {
let request = json!({
"contents": [{
"role": "user",
"parts": [{"text": "Need high effort"}]
}]
});
let converted = build_standard_request_body_with_model_directives(
&request,
"gemini:generate_content",
"gpt-5.4",
"openai",
"openai:chat",
"/v1beta/models/gemini-2.5-pro-high:generateContent",
false,
None,
None,
true,
)
.expect("gemini should convert to openai chat");
assert_eq!(converted["model"], "gpt-5.4");
assert_eq!(converted["reasoning_effort"], "high");
}
#[test]
fn openai_chat_request_uses_typed_canonical_without_changing_target_payloads() {
let request = json!({

18
crates/aether-ai-formats/src/request/standard/mod.rs
View File

@@ -13,8 +13,18 @@ pub use codex::{
CODEX_OPENAI_IMAGE_DEFAULT_VARIATION_PROMPT, CODEX_OPENAI_IMAGE_INTERNAL_MODEL,
};
pub use family::{LocalStandardSourceFamily, LocalStandardSourceMode, LocalStandardSpec};
pub use matrix::{build_standard_request_body, normalize_standard_request_to_openai_chat_request};
pub use normalize::{
build_cross_format_openai_chat_request_body, build_cross_format_openai_responses_request_body,
build_local_openai_chat_request_body, build_local_openai_responses_request_body,
pub use matrix::{
build_standard_request_body, build_standard_request_body_from_canonical_with_model_directives,
build_standard_request_body_with_model_directives,
normalize_standard_request_to_openai_chat_request,
};
pub use normalize::{
build_cross_format_openai_chat_request_body,
build_cross_format_openai_chat_request_body_with_model_directives,
build_cross_format_openai_responses_request_body,
build_cross_format_openai_responses_request_body_with_model_directives,
build_local_openai_chat_request_body,
build_local_openai_chat_request_body_with_model_directives,
build_local_openai_responses_request_body,
build_local_openai_responses_request_body_with_model_directives,
};

235
crates/aether-ai-formats/src/request/standard/normalize.rs
View File

@@ -6,10 +6,26 @@ use aether_ai_formats::protocol::conversion::request::{
use aether_ai_formats::{request_conversion_kind, RequestConversionKind};
use serde_json::{json, Value};
use crate::request::model_directives::apply_model_directive_overrides_from_request;
pub fn build_local_openai_chat_request_body(
body_json: &Value,
mapped_model: &str,
upstream_is_stream: bool,
) -> Option<Value> {
build_local_openai_chat_request_body_with_model_directives(
body_json,
mapped_model,
upstream_is_stream,
false,
)
}
pub fn build_local_openai_chat_request_body_with_model_directives(
body_json: &Value,
mapped_model: &str,
upstream_is_stream: bool,
enable_model_directives: bool,
) -> Option<Value> {
let request_body_object = body_json.as_object()?;
let mut provider_request_body = serde_json::Map::from_iter(
@@ -34,7 +50,14 @@ pub fn build_local_openai_chat_request_body(
}
}
}
Some(Value::Object(provider_request_body))
Some(with_model_directive_overrides(
Value::Object(provider_request_body),
"openai:chat",
mapped_model,
body_json,
None,
enable_model_directives,
))
}
pub fn build_cross_format_openai_chat_request_body(
@@ -42,35 +65,73 @@ pub fn build_cross_format_openai_chat_request_body(
mapped_model: &str,
provider_api_format: &str,
upstream_is_stream: bool,
) -> Option<Value> {
build_cross_format_openai_chat_request_body_with_model_directives(
body_json,
mapped_model,
provider_api_format,
upstream_is_stream,
false,
)
}
pub fn build_cross_format_openai_chat_request_body_with_model_directives(
body_json: &Value,
mapped_model: &str,
provider_api_format: &str,
upstream_is_stream: bool,
enable_model_directives: bool,
) -> Option<Value> {
let conversion_kind = request_conversion_kind("openai:chat", provider_api_format)?;
match conversion_kind {
let provider_request_body = match conversion_kind {
RequestConversionKind::ToClaudeStandard => convert_openai_chat_request_to_claude_request(
body_json,
mapped_model,
upstream_is_stream,
),
)?,
RequestConversionKind::ToGeminiStandard => convert_openai_chat_request_to_gemini_request(
body_json,
mapped_model,
upstream_is_stream,
),
)?,
RequestConversionKind::ToOpenAiResponses => {
convert_openai_chat_request_to_openai_responses_request(
body_json,
mapped_model,
upstream_is_stream,
false,
)
)?
}
_ => None,
}
_ => return None,
};
Some(with_model_directive_overrides(
provider_request_body,
provider_api_format,
mapped_model,
body_json,
None,
enable_model_directives,
))
}
pub fn build_local_openai_responses_request_body(
body_json: &Value,
mapped_model: &str,
require_streaming: bool,
) -> Option<Value> {
build_local_openai_responses_request_body_with_model_directives(
body_json,
mapped_model,
require_streaming,
false,
)
}
pub fn build_local_openai_responses_request_body_with_model_directives(
body_json: &Value,
mapped_model: &str,
require_streaming: bool,
enable_model_directives: bool,
) -> Option<Value> {
let request_body_object = body_json.as_object()?;
let mut provider_request_body = serde_json::Map::from_iter(
@@ -82,7 +143,14 @@ pub fn build_local_openai_responses_request_body(
if require_streaming {
provider_request_body.insert("stream".to_string(), Value::Bool(true));
}
Some(Value::Object(provider_request_body))
Some(with_model_directive_overrides(
Value::Object(provider_request_body),
"openai:responses",
mapped_model,
body_json,
None,
enable_model_directives,
))
}
pub fn build_cross_format_openai_responses_request_body(
@@ -91,41 +159,93 @@ pub fn build_cross_format_openai_responses_request_body(
client_api_format: &str,
provider_api_format: &str,
upstream_is_stream: bool,
) -> Option<Value> {
build_cross_format_openai_responses_request_body_with_model_directives(
body_json,
mapped_model,
client_api_format,
provider_api_format,
upstream_is_stream,
false,
)
}
pub fn build_cross_format_openai_responses_request_body_with_model_directives(
body_json: &Value,
mapped_model: &str,
client_api_format: &str,
provider_api_format: &str,
upstream_is_stream: bool,
enable_model_directives: bool,
) -> Option<Value> {
let chat_like_request = normalize_openai_responses_request_to_openai_chat_request(body_json)?;
let conversion_kind = request_conversion_kind(client_api_format, provider_api_format)?;
match conversion_kind {
RequestConversionKind::ToOpenAIChat => build_local_openai_chat_request_body(
&chat_like_request,
mapped_model,
upstream_is_stream,
),
let provider_request_body = match conversion_kind {
RequestConversionKind::ToOpenAIChat => {
build_local_openai_chat_request_body_with_model_directives(
&chat_like_request,
mapped_model,
upstream_is_stream,
enable_model_directives,
)?
}
RequestConversionKind::ToOpenAiResponses => {
convert_openai_chat_request_to_openai_responses_request(
&chat_like_request,
mapped_model,
upstream_is_stream,
false,
)
)?
}
RequestConversionKind::ToClaudeStandard => convert_openai_chat_request_to_claude_request(
&chat_like_request,
mapped_model,
upstream_is_stream,
),
)?,
RequestConversionKind::ToGeminiStandard => convert_openai_chat_request_to_gemini_request(
&chat_like_request,
mapped_model,
upstream_is_stream,
),
)?,
};
Some(with_model_directive_overrides(
provider_request_body,
provider_api_format,
mapped_model,
body_json,
None,
enable_model_directives,
))
}
fn with_model_directive_overrides(
mut provider_request_body: Value,
provider_api_format: &str,
provider_model: &str,
request_body: &Value,
request_path: Option<&str>,
enable_model_directives: bool,
) -> Value {
if enable_model_directives {
apply_model_directive_overrides_from_request(
&mut provider_request_body,
provider_api_format,
provider_model,
request_body,
request_path,
);
}
provider_request_body
}
#[cfg(test)]
mod tests {
use super::build_local_openai_responses_request_body;
use super::{
build_cross_format_openai_chat_request_body_with_model_directives,
build_cross_format_openai_responses_request_body, build_local_openai_chat_request_body,
build_local_openai_chat_request_body_with_model_directives,
build_local_openai_responses_request_body_with_model_directives,
};
use serde_json::{json, Value};
@@ -204,6 +324,87 @@ mod tests {
);
}
#[test]
fn local_openai_chat_request_body_applies_reasoning_effort_suffix() {
let body_json = json!({
"model": "gpt-5.4-xhigh",
"messages": [{"role": "user", "content": "hello"}],
"reasoning_effort": "low"
});
let provider_request_body = build_local_openai_chat_request_body_with_model_directives(
&body_json,
"gpt-5-upstream",
false,
true,
)
.expect("openai chat body should build");
assert_eq!(provider_request_body["model"], "gpt-5-upstream");
assert_eq!(provider_request_body["reasoning_effort"], "xhigh");
}
#[test]
fn local_openai_chat_request_body_leaves_model_directive_disabled_by_default() {
let body_json = json!({
"model": "gpt-5.4-xhigh",
"messages": [{"role": "user", "content": "hello"}],
"reasoning_effort": "low"
});
let provider_request_body =
build_local_openai_chat_request_body(&body_json, "gpt-5-upstream", false)
.expect("openai chat body should build");
assert_eq!(provider_request_body["model"], "gpt-5-upstream");
assert_eq!(provider_request_body["reasoning_effort"], "low");
}
#[test]
fn local_openai_responses_request_body_applies_reasoning_effort_suffix() {
let body_json = json!({
"model": "gpt-5.4-max",
"input": "hello",
"reasoning": {"effort": "low", "summary": "auto"}
});
let provider_request_body =
build_local_openai_responses_request_body_with_model_directives(
&body_json,
"gpt-5-upstream",
false,
true,
)
.expect("openai responses body should build");
assert_eq!(provider_request_body["model"], "gpt-5-upstream");
assert_eq!(provider_request_body["reasoning"]["summary"], "auto");
assert_eq!(provider_request_body["reasoning"]["effort"], "xhigh");
}
#[test]
fn cross_format_request_body_applies_reasoning_effort_suffix() {
let body_json = json!({
"model": "gpt-5.4-high",
"messages": [{"role": "user", "content": "hello"}],
"reasoning_effort": "low"
});
let provider_request_body =
build_cross_format_openai_chat_request_body_with_model_directives(
&body_json,
"claude-sonnet-4-5",
"claude:messages",
false,
true,
)
.expect("claude body should build");
assert_eq!(provider_request_body["model"], "claude-sonnet-4-5");
assert_eq!(provider_request_body["output_config"]["effort"], "high");
assert_eq!(provider_request_body["thinking"]["budget_tokens"], 4096);
}
#[test]
fn streaming_local_openai_chat_request_body_preserves_stream_options_while_forcing_include_usage(
) {

43
crates/aether-provider-transport/src/same_format_provider.rs
View File

@@ -52,11 +52,14 @@ pub struct SameFormatProviderRequestBehavior {
pub struct SameFormatProviderRequestBodyInput<'a> {
pub body_json: &'a Value,
pub mapped_model: &'a str,
pub provider_api_format: &'a str,
pub source_model: Option<&'a str>,
pub family: SameFormatProviderFamily,
pub body_rules: Option<&'a Value>,
pub upstream_is_stream: bool,
pub kiro_auth_config: Option<&'a KiroAuthConfig>,
pub is_claude_code: bool,
pub enable_model_directives: bool,
}
#[derive(Debug, Clone, Copy)]
@@ -154,6 +157,16 @@ pub fn build_same_format_provider_request_body(
if input.is_claude_code {
crate::claude_code::sanitize_claude_code_request_body(&mut provider_request_body);
}
if input.enable_model_directives {
if let Some(source_model) = input.source_model {
aether_ai_formats::apply_model_directive_overrides_from_model(
&mut provider_request_body,
input.provider_api_format,
input.mapped_model,
source_model,
);
}
}
if !apply_local_body_rules(
&mut provider_request_body,
input.body_rules,
@@ -475,11 +488,14 @@ mod tests {
"messages": [{"role": "user", "content": "hello"}]
}),
mapped_model: "upstream-model",
provider_api_format: "openai:chat",
source_model: Some("client-model"),
family: SameFormatProviderFamily::Standard,
body_rules: None,
upstream_is_stream: true,
kiro_auth_config: None,
is_claude_code: false,
enable_model_directives: false,
})
.expect("body should build");
@@ -487,6 +503,33 @@ mod tests {
assert_eq!(body.get("stream"), Some(&json!(true)));
}
#[test]
fn same_format_body_applies_model_directive_before_body_rules() {
let body = build_same_format_provider_request_body(SameFormatProviderRequestBodyInput {
body_json: &json!({
"model": "gpt-5.4-high",
"messages": [{"role": "user", "content": "hello"}],
"reasoning_effort": "low"
}),
mapped_model: "upstream-model",
provider_api_format: "openai:chat",
source_model: Some("gpt-5.4-high"),
family: SameFormatProviderFamily::Standard,
body_rules: Some(&json!([
{"action":"set","path":"metadata.body_rule_seen","value":true}
])),
upstream_is_stream: false,
kiro_auth_config: None,
is_claude_code: false,
enable_model_directives: true,
})
.expect("body should build");
assert_eq!(body["model"], "upstream-model");
assert_eq!(body["reasoning_effort"], "high");
assert_eq!(body["metadata"]["body_rule_seen"], true);
}
#[test]
fn builds_same_format_headers_with_auth_and_stream_accept() {
let provider_request_body = json!({"model": "upstream-model"});

48
crates/aether-scheduler-core/src/auth.rs
View File

@@ -63,23 +63,44 @@ pub fn auth_constraints_allow_model(
constraints: Option<&SchedulerAuthConstraints>,
requested_model_name: &str,
resolved_global_model_name: &str,
) -> bool {
auth_constraints_allow_model_with_model_directives(
constraints,
requested_model_name,
resolved_global_model_name,
false,
)
}
pub fn auth_constraints_allow_model_with_model_directives(
constraints: Option<&SchedulerAuthConstraints>,
requested_model_name: &str,
resolved_global_model_name: &str,
enable_model_directives: bool,
) -> bool {
let Some(allowed) = constraints.and_then(|constraints| constraints.allowed_models.as_deref())
else {
return true;
};
allowed
.iter()
.any(|value| value == requested_model_name || value == resolved_global_model_name)
let base_model = enable_model_directives
.then(|| aether_ai_formats::model_directive_base_model(requested_model_name))
.flatten();
allowed.iter().any(|value| {
value == requested_model_name
|| value == resolved_global_model_name
|| base_model
.as_ref()
.is_some_and(|base_model| value == base_model)
})
}
#[cfg(test)]
mod tests {
use super::{
api_format_matches_allowed_value, auth_constraints_allow_api_format,
auth_constraints_allow_model, auth_constraints_allow_provider,
provider_matches_allowed_value, SchedulerAuthConstraints,
auth_constraints_allow_model, auth_constraints_allow_model_with_model_directives,
auth_constraints_allow_provider, provider_matches_allowed_value, SchedulerAuthConstraints,
};
fn sample_constraints() -> SchedulerAuthConstraints {
@@ -200,6 +221,23 @@ mod tests {
));
}
#[test]
fn model_directive_base_model_requires_explicit_enablement() {
let constraints = sample_constraints();
assert!(!auth_constraints_allow_model(
Some(&constraints),
"gpt-5-high",
"gpt-5-high"
));
assert!(auth_constraints_allow_model_with_model_directives(
Some(&constraints),
"gpt-5-high",
"gpt-5-high",
true
));
}
#[test]
fn api_format_allowed_value_matches_current_signatures_only() {
assert!(api_format_matches_allowed_value(

24
crates/aether-scheduler-core/src/candidate/enumeration.rs
View File

@@ -9,6 +9,20 @@ use super::types::{
pub fn enumerate_minimal_candidate_selection(
input: EnumerateMinimalCandidateSelectionInput<'_>,
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, DataLayerError> {
enumerate_minimal_candidate_selection_inner(input, false)
}
pub fn enumerate_minimal_candidate_selection_with_model_directives(
input: EnumerateMinimalCandidateSelectionInput<'_>,
enable_model_directives: bool,
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, DataLayerError> {
enumerate_minimal_candidate_selection_inner(input, enable_model_directives)
}
fn enumerate_minimal_candidate_selection_inner(
input: EnumerateMinimalCandidateSelectionInput<'_>,
enable_model_directives: bool,
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, DataLayerError> {
let EnumerateMinimalCandidateSelectionInput {
rows,
@@ -26,10 +40,11 @@ pub fn enumerate_minimal_candidate_selection(
if !crate::auth_constraints_allow_api_format(auth_constraints, normalized_api_format) {
return Ok(Vec::new());
}
if !crate::auth_constraints_allow_model(
if !crate::auth_constraints_allow_model_with_model_directives(
auth_constraints,
requested_model_name,
resolved_global_model_name,
enable_model_directives,
) {
return Ok(Vec::new());
}
@@ -48,7 +63,12 @@ pub fn enumerate_minimal_candidate_selection(
continue;
}
let Some((selected_provider_model_name, mapping_matched_model)) =
crate::resolve_provider_model_name(&row, requested_model_name, normalized_api_format)
crate::resolve_provider_model_name_with_model_directives(
&row,
requested_model_name,
normalized_api_format,
enable_model_directives,
)
else {
continue;
};

1
crates/aether-scheduler-core/src/candidate/mod.rs
View File

@@ -8,6 +8,7 @@ pub use capability::{
};
pub use enumeration::{
collect_global_model_names_for_required_capability, enumerate_minimal_candidate_selection,
enumerate_minimal_candidate_selection_with_model_directives,
};
pub use selectability::{
auth_api_key_concurrency_limit_reached, candidate_is_selectable_with_runtime_state,

12
crates/aether-scheduler-core/src/lib.rs
View File

@@ -13,13 +13,14 @@ pub use affinity::{
};
pub use auth::{
api_format_matches_allowed_value, auth_constraints_allow_api_format,
auth_constraints_allow_model, auth_constraints_allow_provider, provider_matches_allowed_value,
SchedulerAuthConstraints,
auth_constraints_allow_model, auth_constraints_allow_model_with_model_directives,
auth_constraints_allow_provider, provider_matches_allowed_value, SchedulerAuthConstraints,
};
pub use candidate::{
auth_api_key_concurrency_limit_reached, candidate_is_selectable_with_runtime_state,
candidate_runtime_skip_reason_with_state, candidate_supports_required_capability,
collect_global_model_names_for_required_capability, enumerate_minimal_candidate_selection,
enumerate_minimal_candidate_selection_with_model_directives,
requested_capability_priority_for_candidate, CandidateRuntimeSelectabilityInput,
EnumerateMinimalCandidateSelectionInput, SchedulerMinimalCandidateSelectionCandidate,
SchedulerPriorityMode,
@@ -35,8 +36,11 @@ pub use health::{
};
pub use model::{
candidate_model_names, extract_global_priority_for_format, matches_model_mapping,
normalize_api_format, resolve_provider_model_name, resolve_requested_global_model_name,
row_supports_requested_model, row_supports_required_capability, select_provider_model_name,
normalize_api_format, resolve_provider_model_name,
resolve_provider_model_name_with_model_directives, resolve_requested_global_model_name,
resolve_requested_global_model_name_with_model_directives, row_supports_requested_model,
row_supports_requested_model_with_model_directives, row_supports_required_capability,
select_provider_model_name,
};
pub use provider::{build_provider_concurrent_limit_map, should_skip_provider_quota};
pub use ranking::{

226
crates/aether-scheduler-core/src/model.rs
View File

@@ -1,3 +1,4 @@
use std::borrow::Cow;
use std::collections::BTreeSet;
use aether_data_contracts::repository::candidate_selection::{
@@ -11,39 +12,79 @@ pub fn resolve_requested_global_model_name(
requested_model_name: &str,
api_format: &str,
) -> Option<String> {
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
})
})
.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, api_format)
&& mapping.name == requested_model_name
resolve_requested_global_model_name_with_model_directives(
rows,
requested_model_name,
api_format,
false,
)
}
pub fn resolve_requested_global_model_name_with_model_directives(
rows: &[StoredMinimalCandidateSelectionRow],
requested_model_name: &str,
api_format: &str,
enable_model_directives: bool,
) -> Option<String> {
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
})
})
.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, 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| {
patterns
.iter()
.any(|pattern| matches_model_mapping(pattern, requested_model_name))
})
})
})
})
.or_else(|| {
resolve_global_model_name_by(rows, |row| {
row.global_model_mappings.as_ref().is_some_and(|patterns| {
patterns
.iter()
.any(|pattern| matches_model_mapping(pattern, requested_model_name))
})
})
})
},
)
}
pub fn row_supports_requested_model(
row: &StoredMinimalCandidateSelectionRow,
requested_model_name: &str,
api_format: &str,
) -> bool {
row_supports_requested_model_with_model_directives(row, requested_model_name, api_format, false)
}
pub fn row_supports_requested_model_with_model_directives(
row: &StoredMinimalCandidateSelectionRow,
requested_model_name: &str,
api_format: &str,
enable_model_directives: bool,
) -> bool {
requested_model_name_candidates(requested_model_name, enable_model_directives).any(
|requested_model_name| {
row_supports_requested_model_exact(row, requested_model_name.as_ref(), api_format)
},
)
}
fn row_supports_requested_model_exact(
row: &StoredMinimalCandidateSelectionRow,
requested_model_name: &str,
api_format: &str,
) -> bool {
row.global_model_name == requested_model_name
|| row.model_provider_model_name == requested_model_name
@@ -87,6 +128,15 @@ pub fn resolve_provider_model_name(
row: &StoredMinimalCandidateSelectionRow,
requested_model_name: &str,
api_format: &str,
) -> Option<(String, Option<String>)> {
resolve_provider_model_name_with_model_directives(row, requested_model_name, api_format, false)
}
pub fn resolve_provider_model_name_with_model_directives(
row: &StoredMinimalCandidateSelectionRow,
requested_model_name: &str,
api_format: &str,
enable_model_directives: bool,
) -> Option<(String, Option<String>)> {
let selected_provider_model_name = select_provider_model_name(row, api_format);
let Some(key_allowed_models) = row.key_allowed_models.as_ref() else {
@@ -103,6 +153,16 @@ pub fn resolve_provider_model_name(
return Some((selected_provider_model_name, None));
}
if enable_model_directives {
if let Some(base_model) =
aether_ai_formats::model_directive_base_model(requested_model_name)
{
if key_allowed_models.iter().any(|value| value == &base_model) {
return Some((selected_provider_model_name, Some(base_model)));
}
}
}
let mut sorted_allowed_models = key_allowed_models
.iter()
.map(String::as_str)
@@ -302,9 +362,26 @@ fn api_format_matches(left: &str, right: &str) -> bool {
normalize_api_format(left) == normalize_api_format(right)
}
fn requested_model_name_candidates(
requested_model_name: &str,
enable_model_directives: bool,
) -> impl Iterator<Item = Cow<'_, str>> {
let requested_model_name = requested_model_name.trim();
let base_model = enable_model_directives
.then(|| aether_ai_formats::model_directive_base_model(requested_model_name))
.flatten();
std::iter::once(Cow::Borrowed(requested_model_name)).chain(base_model.map(Cow::Owned))
}
#[cfg(test)]
mod tests {
use super::matches_model_mapping;
use super::{
matches_model_mapping, resolve_provider_model_name,
resolve_provider_model_name_with_model_directives,
resolve_requested_global_model_name_with_model_directives, row_supports_requested_model,
row_supports_requested_model_with_model_directives,
};
use aether_data_contracts::repository::candidate_selection::StoredMinimalCandidateSelectionRow;
#[test]
fn model_mapping_match_is_case_insensitive() {
@@ -322,4 +399,103 @@ mod tests {
fn invalid_model_mapping_pattern_returns_false() {
assert!(!matches_model_mapping("([a-z", "gpt-4o"));
}
#[test]
fn model_directive_suffix_matches_base_model_as_fallback() {
let row = sample_row("gpt-5.4", "gpt-5.4-upstream");
assert!(!row_supports_requested_model(
&row,
"gpt-5.4-xhigh",
"openai:chat"
));
assert!(row_supports_requested_model_with_model_directives(
&row,
"gpt-5.4-xhigh",
"openai:chat",
true
));
assert_eq!(
resolve_requested_global_model_name_with_model_directives(
&[row],
"gpt-5.4-xhigh",
"openai:chat",
true
)
.as_deref(),
Some("gpt-5.4")
);
}
#[test]
fn model_directive_suffix_prefers_exact_model_before_base_fallback() {
let exact = sample_row("gpt-5.4-high", "gpt-5.4-high-upstream");
let base = sample_row("gpt-5.4", "gpt-5.4-upstream");
assert_eq!(
resolve_requested_global_model_name_with_model_directives(
&[base, exact],
"gpt-5.4-high",
"openai:chat",
true
)
.as_deref(),
Some("gpt-5.4-high")
);
}
#[test]
fn model_directive_base_model_satisfies_key_allowed_models() {
let mut row = sample_row("gpt-5.4", "gpt-5.4-upstream");
row.key_allowed_models = Some(vec!["gpt-5.4".to_string()]);
assert!(resolve_provider_model_name(&row, "gpt-5.4-max", "openai:chat").is_none());
let resolved = resolve_provider_model_name_with_model_directives(
&row,
"gpt-5.4-max",
"openai:chat",
true,
)
.expect("base model should satisfy key allowed models");
assert_eq!(resolved.0, "gpt-5.4-upstream");
assert_eq!(resolved.1.as_deref(), Some("gpt-5.4"));
}
fn sample_row(
global_model_name: &str,
model_provider_model_name: &str,
) -> StoredMinimalCandidateSelectionRow {
StoredMinimalCandidateSelectionRow {
provider_id: "provider-1".to_string(),
provider_name: "Provider".to_string(),
provider_type: "openai".to_string(),
provider_priority: 0,
provider_is_active: true,
endpoint_id: "endpoint-1".to_string(),
endpoint_api_format: "openai:chat".to_string(),
endpoint_api_family: None,
endpoint_kind: None,
endpoint_is_active: true,
key_id: "key-1".to_string(),
key_name: "Key".to_string(),
key_auth_type: "api_key".to_string(),
key_is_active: true,
key_api_formats: None,
key_allowed_models: None,
key_capabilities: None,
key_internal_priority: 0,
key_global_priority_by_format: None,
model_id: format!("model-{global_model_name}"),
global_model_id: format!("global-{global_model_name}"),
global_model_name: global_model_name.to_string(),
global_model_mappings: None,
global_model_supports_streaming: Some(true),
model_provider_model_name: model_provider_model_name.to_string(),
model_provider_model_mappings: None,
model_supports_streaming: Some(true),
model_is_active: true,
model_is_available: true,
}
}
}