Aether/crates/aether-scheduler-core/src/model.rs

use std::borrow::Cow;
use std::collections::BTreeSet;

use aether_data_contracts::repository::candidate_selection::{
    StoredMinimalCandidateSelectionRow, StoredProviderModelMapping,
};
use aether_data_contracts::DataLayerError;
use regex::RegexBuilder;

pub fn resolve_requested_global_model_name(
    rows: &[StoredMinimalCandidateSelectionRow],
    requested_model_name: &str,
    api_format: &str,
) -> Option<String> {
    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, 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| {
                            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
        || 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
                })
            })
        || 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>(
    rows: &[StoredMinimalCandidateSelectionRow],
    matches: F,
) -> Option<String>
where
    F: Fn(&StoredMinimalCandidateSelectionRow) -> bool,
{
    let mut best_match = None::<&str>;
    for row in rows.iter().filter(|row| matches(row)) {
        let candidate = row.global_model_name.trim();
        if candidate.is_empty() {
            continue;
        }
        if best_match.is_none_or(|current| candidate < current) {
            best_match = Some(candidate);
        }
    }
    best_match.map(ToOwned::to_owned)
}

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 {
        return Some((selected_provider_model_name, None));
    };
    if key_allowed_models.is_empty() {
        return None;
    }

    if key_allowed_models
        .iter()
        .any(|value| value == requested_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)
        .map(str::trim)
        .filter(|value| !value.is_empty())
        .collect::<Vec<_>>();
    sorted_allowed_models.sort_unstable();

    for &allowed_model in &sorted_allowed_models {
        if row_has_candidate_model_name(row, api_format, allowed_model) {
            let allowed_model = allowed_model.to_owned();
            return Some((selected_provider_model_name.clone(), Some(allowed_model)));
        }
    }

    let global_model_mappings = row.global_model_mappings.as_ref()?;
    for &allowed_model in &sorted_allowed_models {
        for pattern in global_model_mappings {
            if matches_model_mapping(pattern, allowed_model) {
                let allowed_model = allowed_model.to_owned();
                // Regex mappings prove the key can serve this global model; they do not
                // override the provider model chosen from the Provider model mapping.
                return Some((selected_provider_model_name.clone(), Some(allowed_model)));
            }
        }
    }

    None
}

pub fn select_provider_model_name(
    row: &StoredMinimalCandidateSelectionRow,
    api_format: &str,
) -> String {
    let Some(mappings) = row.model_provider_model_mappings.as_ref() else {
        return row.model_provider_model_name.clone();
    };

    mappings
        .iter()
        .filter(|mapping| mapping_scope_matches(mapping, row, api_format))
        .min_by(|left, right| {
            left.priority
                .cmp(&right.priority)
                .then(left.name.cmp(&right.name))
        })
        .map(|mapping| mapping.name.clone())
        .unwrap_or_else(|| 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()]);
    if let Some(mappings) = row.model_provider_model_mappings.as_ref() {
        for mapping in mappings {
            if mapping_scope_matches(mapping, row, api_format) {
                names.insert(mapping.name.clone());
            }
        }
    }
    names
}

fn mapping_scope_matches(
    mapping: &StoredProviderModelMapping,
    row: &StoredMinimalCandidateSelectionRow,
    api_format: &str,
) -> bool {
    let api_format_matches_scope = mapping.api_formats.as_ref().is_none_or(|api_formats| {
        api_formats
            .iter()
            .any(|value| api_format_matches(value, api_format))
    });
    if !api_format_matches_scope {
        return false;
    }

    mapping.endpoint_ids.as_ref().is_none_or(|endpoint_ids| {
        endpoint_ids
            .iter()
            .any(|endpoint_id| endpoint_id == &row.endpoint_id)
    })
}

pub fn row_supports_required_capability(
    row: &StoredMinimalCandidateSelectionRow,
    required_capability: &str,
) -> bool {
    capabilities_support_required_capability(row.key_capabilities.as_ref(), required_capability)
}

fn capabilities_support_required_capability(
    capabilities: Option<&serde_json::Value>,
    required_capability: &str,
) -> bool {
    let required_capability = required_capability.trim();
    if required_capability.is_empty() {
        return true;
    }
    let Some(capabilities) = capabilities else {
        return false;
    };

    if let Some(object) = capabilities.as_object() {
        return object.iter().any(|(key, value)| {
            key.eq_ignore_ascii_case(required_capability)
                && match value {
                    serde_json::Value::Bool(value) => *value,
                    serde_json::Value::String(value) => value.eq_ignore_ascii_case("true"),
                    serde_json::Value::Number(value) => {
                        value.as_i64().is_some_and(|value| value > 0)
                    }
                    _ => false,
                }
        });
    }

    if let Some(items) = capabilities.as_array() {
        return items.iter().any(|value| {
            value
                .as_str()
                .is_some_and(|value| value.eq_ignore_ascii_case(required_capability))
        });
    }

    false
}

pub fn matches_model_mapping(pattern: &str, model_name: &str) -> bool {
    if pattern.eq_ignore_ascii_case(model_name) {
        return true;
    }

    let regex_pattern = format!("^(?:{pattern})$");
    let Ok(compiled) = RegexBuilder::new(&regex_pattern)
        .case_insensitive(true)
        .build()
    else {
        return false;
    };
    compiled.is_match(model_name)
}

pub fn extract_global_priority_for_format(
    raw: Option<&serde_json::Value>,
    api_format: &str,
) -> Result<Option<i32>, DataLayerError> {
    let Some(raw) = raw else {
        return Ok(None);
    };
    let Some(object) = raw.as_object() else {
        return Err(DataLayerError::UnexpectedValue(
            "provider_api_keys.global_priority_by_format is not a JSON object".to_string(),
        ));
    };

    let Some(value) = object
        .iter()
        .find(|(key, _)| api_format_matches(key, api_format))
        .map(|(_, value)| value)
    else {
        return Ok(None);
    };

    if let Some(value) = value.as_i64() {
        return i32::try_from(value).map(Some).map_err(|_| {
            DataLayerError::UnexpectedValue(format!(
                "invalid provider_api_keys.global_priority_by_format value: {value}"
            ))
        });
    }

    if let Some(value) = value.as_str() {
        let value = value.trim().parse::<i32>().map_err(|_| {
            DataLayerError::UnexpectedValue(format!(
                "invalid provider_api_keys.global_priority_by_format value: {value}"
            ))
        })?;
        return Ok(Some(value));
    }

    Err(DataLayerError::UnexpectedValue(
        "provider_api_keys.global_priority_by_format contains a non-integer value".to_string(),
    ))
}

pub fn normalize_api_format(value: &str) -> String {
    aether_ai_formats::normalize_api_format_alias(value)
}

fn row_has_candidate_model_name(
    row: &StoredMinimalCandidateSelectionRow,
    api_format: &str,
    model_name: &str,
) -> bool {
    row.model_provider_model_name == model_name
        || row
            .model_provider_model_mappings
            .as_ref()
            .is_some_and(|mappings| {
                mappings.iter().any(|mapping| {
                    mapping_scope_matches(mapping, row, api_format) && mapping.name == model_name
                })
            })
}

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, 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, StoredProviderModelMapping,
    };

    #[test]
    fn model_mapping_match_is_case_insensitive() {
        assert!(matches_model_mapping("gpt-4o", "GPT-4O"));
        assert!(matches_model_mapping("gpt-5(?:\\.\\d+)?", "GPT-5.1"));
    }

    #[test]
    fn model_mapping_match_is_anchored_to_full_text() {
        assert!(matches_model_mapping("gpt-4o", "gpt-4o"));
        assert!(!matches_model_mapping("gpt-4o", "gpt-4o-mini"));
    }

    #[test]
    fn invalid_model_mapping_pattern_returns_false() {
        assert!(!matches_model_mapping("([a-z", "gpt-4o"));
    }

    #[test]
    fn regex_allowed_model_does_not_replace_selected_provider_model_name() {
        let mut row = sample_row("gpt-5", "gpt-5-upstream");
        row.key_allowed_models = Some(vec!["gpt-5.4".to_string()]);
        row.global_model_mappings = Some(vec!["gpt-5(?:\\.\\d+)?".to_string()]);
        row.model_provider_model_mappings = Some(vec![StoredProviderModelMapping {
            name: "gpt-5-canonical-upstream".to_string(),
            priority: 1,
            api_formats: Some(vec!["openai:chat".to_string()]),
            endpoint_ids: None,
        }]);

        let resolved = resolve_provider_model_name(&row, "gpt-5", "openai:chat")
            .expect("regex-matched allowed model should allow the key");

        assert_eq!(resolved.0, "gpt-5-canonical-upstream");
        assert_eq!(resolved.1.as_deref(), Some("gpt-5.4"));
    }

    #[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,
        }
    }
}