use std::fmt;
use super::discovery::{EndpointInfo, Pricing, parse_price};
pub(crate) const VALIDATED_MIX_CACHED: f64 = 0.977;
pub(crate) const VALIDATED_MIX_INPUT: f64 = 0.014;
pub(crate) const VALIDATED_MIX_OUTPUT: f64 = 0.009;
#[must_use]
#[expect(clippy::cast_precision_loss)]
pub(crate) fn estimate_cost(
price: &Pricing,
total_tokens: f64,
requests: usize,
flags: &mut Vec<String>,
) -> f64 {
let cache_read = if let Some(p) = price.input_cache_read.as_deref().and_then(parse_price) {
p
} else {
flags
.push("no cache-read price advertised; estimate assumes full prompt price".to_string());
price.prompt.as_deref().and_then(parse_price).unwrap_or(0.0)
};
let prompt = price.prompt.as_deref().and_then(parse_price).unwrap_or(0.0);
let completion = price
.completion
.as_deref()
.and_then(parse_price)
.unwrap_or(0.0);
let request = price
.request
.as_deref()
.and_then(parse_price)
.unwrap_or(0.0);
let blended = VALIDATED_MIX_CACHED * cache_read
+ VALIDATED_MIX_INPUT * prompt
+ VALIDATED_MIX_OUTPUT * completion;
blended * total_tokens + request * requests as f64
}
pub(crate) struct SelectionInput {
pub endpoint: EndpointInfo,
pub est_cost: f64,
}
#[derive(Debug, Clone, PartialEq)]
pub(crate) enum ExclusionReason {
Status(String),
FreeVariant,
ContextTooSmall(i64),
ContextUnknown,
NotInAllowlist,
Outlier(f64),
NotSelected(f64),
Padding,
}
impl fmt::Display for ExclusionReason {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Status(s) => write!(f, "status {s}"),
Self::FreeVariant => write!(f, "free variant (:free tag)"),
Self::ContextTooSmall(cl) => write!(f, "context too small ({cl} tokens)"),
Self::ContextUnknown => write!(f, "context unknown"),
Self::NotInAllowlist => write!(f, "not in provider allowlist"),
Self::Outlier(est) => write!(f, "cost outlier (est ${est:.4} > 3×median)"),
Self::NotSelected(est) => write!(f, "not selected (est ${est:.4})"),
Self::Padding => write!(f, "padding (expected to fail)"),
}
}
}
pub(crate) struct SelectionDecision {
pub selected: bool,
pub reason: Option<ExclusionReason>,
}
impl SelectionDecision {
#[must_use]
pub(crate) fn reason_text(&self) -> String {
match (&self.reason, self.selected) {
(Some(r), _) => r.to_string(),
(None, true) => "selected".to_string(),
(None, false) => "not selected".to_string(),
}
}
#[must_use]
pub(crate) fn is_healthy(&self) -> bool {
matches!(
self.reason,
None | Some(ExclusionReason::NotSelected(_) | ExclusionReason::Outlier(_))
)
}
}
#[must_use]
pub(crate) fn is_healthy_status(s: Option<&str>) -> bool {
s == Some("0")
}
#[must_use]
pub(crate) fn is_free_variant(tag: &str) -> bool {
tag.ends_with(":free")
}
#[must_use]
pub(crate) fn classify_endpoint(
endpoint: &EndpointInfo,
min_context: i64,
allowlist: Option<&[String]>,
) -> (bool, Option<ExclusionReason>) {
let in_allowlist = allowlist.is_none_or(|wl| wl.iter().any(|t| t == &endpoint.tag));
let reason = if !in_allowlist {
Some(ExclusionReason::NotInAllowlist)
} else if is_free_variant(&endpoint.tag) {
Some(ExclusionReason::FreeVariant)
} else {
match endpoint.context_length {
None => Some(ExclusionReason::ContextUnknown),
Some(cl) if cl < min_context => Some(ExclusionReason::ContextTooSmall(cl)),
Some(_) if !is_healthy_status(endpoint.status.as_deref()) => Some(
ExclusionReason::Status(endpoint.status.clone().unwrap_or_default()),
),
Some(_) => None,
}
};
(reason.is_none(), reason)
}
#[must_use]
pub(crate) fn select_providers(
input: &[SelectionInput],
min_context: i64,
allowlist: Option<&[String]>,
) -> Vec<SelectionDecision> {
let candidates = candidate_indices(input, allowlist);
if allowlist.is_some() && !candidates.is_empty() && candidates.len() <= 2 {
let mut decisions = base_decisions(input, min_context, allowlist);
for &idx in &candidates {
decisions[idx].selected = true;
}
return decisions;
}
let mut decisions = base_decisions(input, min_context, allowlist);
let mut healthy: Vec<usize> = candidates
.iter()
.copied()
.filter(|&idx| decisions[idx].reason.is_none())
.collect();
if healthy.is_empty() {
return decisions; }
if healthy.len() < 3 {
pad_to_three(&mut decisions, &candidates, input);
return decisions;
}
healthy.sort_by(|&a, &b| input[a].est_cost.total_cmp(&input[b].est_cost));
let median = median_of(&healthy, input);
let outliers: Vec<usize> = healthy
.iter()
.copied()
.filter(|&idx| input[idx].est_cost > 3.0 * median)
.collect();
let mut pool: Vec<usize> = healthy
.iter()
.copied()
.filter(|&idx| input[idx].est_cost <= 3.0 * median)
.collect();
if pool.len() < 3 {
let mut by_cost = outliers.clone();
by_cost.sort_by(|&a, &b| input[a].est_cost.total_cmp(&input[b].est_cost));
for idx in by_cost {
if pool.len() >= 3 {
break;
}
pool.push(idx);
}
}
let n = selection_target(healthy.len());
for &idx in pool.iter().take(n) {
decisions[idx] = SelectionDecision {
selected: true,
reason: None,
};
}
for &idx in pool.iter().skip(n) {
decisions[idx] = SelectionDecision {
selected: false,
reason: Some(ExclusionReason::NotSelected(input[idx].est_cost)),
};
}
for &idx in &outliers {
if !pool.contains(&idx) {
decisions[idx] = SelectionDecision {
selected: false,
reason: Some(ExclusionReason::Outlier(input[idx].est_cost)),
};
}
}
decisions
}
fn candidate_indices(input: &[SelectionInput], allowlist: Option<&[String]>) -> Vec<usize> {
match allowlist {
Some(wl) if !wl.is_empty() => input
.iter()
.enumerate()
.filter(|(_, i)| wl.iter().any(|t| t == &i.endpoint.tag))
.map(|(idx, _)| idx)
.collect(),
Some(_) => Vec::new(),
None => (0..input.len()).collect(),
}
}
fn base_decisions(
input: &[SelectionInput],
min_context: i64,
allowlist: Option<&[String]>,
) -> Vec<SelectionDecision> {
input
.iter()
.map(|i| {
let (_, reason) = classify_endpoint(&i.endpoint, min_context, allowlist);
SelectionDecision {
selected: false,
reason,
}
})
.collect()
}
fn pad_to_three(
decisions: &mut [SelectionDecision],
candidates: &[usize],
input: &[SelectionInput],
) {
for &idx in candidates {
if decisions[idx].reason.is_none() {
decisions[idx].selected = true;
}
}
let mut by_cost: Vec<usize> = candidates
.iter()
.copied()
.filter(|&idx| !decisions[idx].selected)
.collect();
by_cost.sort_by(|&a, &b| input[a].est_cost.total_cmp(&input[b].est_cost));
for idx in by_cost {
if decisions.iter().filter(|d| d.selected).count() >= 3 {
break;
}
decisions[idx] = SelectionDecision {
selected: true,
reason: Some(ExclusionReason::Padding),
};
}
}
#[must_use]
#[expect(
clippy::cast_possible_truncation,
clippy::cast_precision_loss,
clippy::cast_sign_loss
)]
pub(crate) fn selection_target(healthy_count: usize) -> usize {
(3usize).max((healthy_count as f64 * 0.8).ceil() as usize)
}
#[must_use]
pub(crate) fn effective_target_count(
healthy_count: usize,
allowlist_matches: Option<usize>,
) -> usize {
match (healthy_count, allowlist_matches) {
(0, _) => 0,
(_, Some(m)) if m > 0 && m <= 2 => m,
_ => selection_target(healthy_count),
}
}
fn median_of(indices: &[usize], input: &[SelectionInput]) -> f64 {
let mut costs: Vec<f64> = indices.iter().map(|&i| input[i].est_cost).collect();
costs.sort_by(f64::total_cmp);
let mid = costs.len() / 2;
if costs.len().is_multiple_of(2) {
f64::midpoint(costs[mid - 1], costs[mid])
} else {
costs[mid]
}
}
#[must_use]
#[expect(clippy::cast_precision_loss)]
pub(crate) fn plan_cost(
selected: &[SelectionDecision],
inputs: &[SelectionInput],
cap: f64,
) -> (f64, f64) {
let selected_count = selected.iter().filter(|d| d.selected).count();
let total: f64 = selected
.iter()
.zip(inputs)
.filter(|(d, _)| d.selected)
.map(|(_, i)| i.est_cost)
.sum();
let guard = cap * 2.0 / (selected_count.max(1) as f64);
(total, guard)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::bench_openrouter::discovery::Pricing;
fn ep(tag: &str, context: i64, status: &str) -> EndpointInfo {
ep_ctx(tag, Some(context), status)
}
fn ep_ctx(tag: &str, context: Option<i64>, status: &str) -> EndpointInfo {
EndpointInfo {
tag: tag.to_string(),
name: tag.to_string(),
provider_name: tag.to_string(),
context_length: context,
quantization: None,
status: Some(status.to_string()),
supports_implicit_caching: Some(true),
pricing: Some(Pricing {
prompt: Some("0.000001".to_string()),
completion: Some("0.000002".to_string()),
request: Some("0".to_string()),
input_cache_read: Some("0.0000001".to_string()),
}),
}
}
fn input(ep: EndpointInfo, est: f64) -> SelectionInput {
SelectionInput {
endpoint: ep,
est_cost: est,
}
}
#[test]
fn healthy_status_parsing() {
assert!(is_healthy_status(Some("0")));
for s in ["-1", "-2", "-3", "-5", "-10", "", "abc"] {
assert!(!is_healthy_status(Some(s)), "status {s:?}");
}
assert!(!is_healthy_status(None));
}
#[test]
fn free_variant_detection() {
assert!(is_free_variant("groq/llama-3.1-70b:free"));
assert!(!is_free_variant("groq/llama-3.1-70b"));
assert!(!is_free_variant(""));
}
#[test]
fn estimate_cost_uses_cache_read_and_request_fees() {
let price = Pricing {
prompt: Some("0.000002".to_string()),
completion: Some("0.000008".to_string()),
request: Some("0.0005".to_string()),
input_cache_read: Some("0.0000002".to_string()),
};
let mut flags = Vec::new();
let est = estimate_cost(&price, 1000.0, 4, &mut flags);
assert!(flags.is_empty(), "cache-read price present → no flag");
let expected =
(0.977 * 0.000_000_2 + 0.014 * 0.000_002 + 0.009 * 0.000_008) * 1000.0 + 0.0005 * 4.0;
assert!((est - expected).abs() < 1e-12);
}
#[test]
fn estimate_cost_falls_back_to_prompt_with_flag() {
let price = Pricing {
prompt: Some("0.000002".to_string()),
completion: Some("0.000008".to_string()),
request: Some("0".to_string()),
input_cache_read: None, };
let mut flags = Vec::new();
let est = estimate_cost(&price, 1000.0, 1, &mut flags);
assert_eq!(flags.len(), 1);
assert!(flags[0].contains("no cache-read price advertised"));
let expected = (0.977 + 0.014) * 0.000_002 * 1000.0 + 0.009 * 0.000_008 * 1000.0;
assert!((est - expected).abs() < 1e-12);
}
#[test]
fn selection_target_count_and_outlier() {
let mut inputs = Vec::new();
for i in 0..8 {
let est = if i == 7 { 100.0 } else { f64::from(i + 1) };
inputs.push(input(ep(&format!("p{i}"), 200_000, "0"), est));
}
let decisions = select_providers(&inputs, 128_000, None);
let selected: Vec<_> = decisions.iter().filter(|d| d.selected).collect();
assert_eq!(selected.len(), 7);
assert_eq!(decisions[7].reason, Some(ExclusionReason::Outlier(100.0)));
for d in &decisions[..6] {
assert!(d.selected);
}
}
#[test]
fn selection_pads_healthy_lt_3() {
let mut inputs = Vec::new();
let h0 = input(ep("h0", 200_000, "0"), 1.0);
let h1 = input(ep("h1", 200_000, "0"), 2.0);
let bad = input(ep("bad", 200_000, "-10"), 3.0);
inputs.extend([h0, h1, bad]);
let decisions = select_providers(&inputs, 128_000, None);
let selected: Vec<_> = decisions.iter().filter(|d| d.selected).collect();
assert_eq!(selected.len(), 3);
assert!(decisions[0].selected && decisions[1].selected);
assert!(decisions[2].selected);
assert_eq!(decisions[2].reason, Some(ExclusionReason::Padding));
}
#[test]
fn selection_zero_healthy_is_empty() {
let inputs = vec![
input(ep("a", 200_000, "-10"), 1.0),
input(ep("b", 200_000, "-2"), 2.0),
];
let decisions = select_providers(&inputs, 128_000, None);
assert!(decisions.iter().all(|d| !d.selected));
}
#[test]
fn selection_allowlist_restricts_and_short_allowlist_wins() {
let mut inputs = Vec::new();
for i in 0..4 {
inputs.push(input(ep(&format!("p{i}"), 200_000, "0"), f64::from(i + 1)));
}
let wl = vec!["p1".to_string(), "p3".to_string()];
let decisions = select_providers(&inputs, 128_000, Some(&wl));
assert_eq!(decisions.iter().filter(|d| d.selected).count(), 2);
assert!(decisions[1].selected && decisions[3].selected);
assert!(!decisions[0].selected && !decisions[2].selected);
let wl = vec!["p0".to_string()];
let decisions = select_providers(&inputs, 128_000, Some(&wl));
assert_eq!(decisions.iter().filter(|d| d.selected).count(), 1);
assert!(decisions[0].selected);
}
#[test]
fn selection_min_context_excludes_small_contexts() {
let inputs = vec![
input(ep("small", 32_000, "0"), 5.0),
input(ep_ctx("none", None, "0"), 5.0),
input(ep("ok1", 200_000, "0"), 1.0),
input(ep("ok2", 200_000, "0"), 2.0),
input(ep("ok3", 200_000, "0"), 3.0),
input(ep("ok4", 200_000, "0"), 4.0),
];
let decisions = select_providers(&inputs, 128_000, None);
assert_eq!(
decisions[0].reason,
Some(ExclusionReason::ContextTooSmall(32_000))
);
assert_eq!(decisions[1].reason, Some(ExclusionReason::ContextUnknown));
assert!(!decisions[0].selected && !decisions[1].selected);
assert_eq!(decisions.iter().filter(|d| d.selected).count(), 4);
for (i, d) in decisions.iter().enumerate().take(6).skip(2) {
assert!(d.selected, "endpoint {i} should be selected");
}
}
#[test]
fn plan_cost_sums_selected_and_computes_guard() {
let inputs = vec![
input(ep("a", 200_000, "0"), 0.1),
input(ep("b", 200_000, "0"), 0.2),
input(ep("c", 200_000, "-10"), 0.3),
];
let decisions = select_providers(&inputs, 128_000, None);
let (total, guard) = plan_cost(&decisions, &inputs, 2.0);
assert!((total - 0.6).abs() < 1e-12);
assert!((guard - 4.0 / 3.0).abs() < 1e-12);
}
#[test]
fn exclusion_reason_displays() {
assert_eq!(
ExclusionReason::Status("-10".to_string()).to_string(),
"status -10"
);
assert_eq!(
ExclusionReason::FreeVariant.to_string(),
"free variant (:free tag)"
);
assert_eq!(
ExclusionReason::ContextTooSmall(32_000).to_string(),
"context too small (32000 tokens)"
);
assert_eq!(
ExclusionReason::ContextUnknown.to_string(),
"context unknown"
);
assert_eq!(
ExclusionReason::NotInAllowlist.to_string(),
"not in provider allowlist"
);
assert_eq!(
ExclusionReason::Padding.to_string(),
"padding (expected to fail)"
);
}
#[test]
fn selection_decision_reason_text() {
let d = SelectionDecision {
selected: false,
reason: Some(ExclusionReason::Padding),
};
assert_eq!(d.reason_text(), "padding (expected to fail)");
let selected = SelectionDecision {
selected: true,
reason: None,
};
assert_eq!(selected.reason_text(), "selected");
let unselected = SelectionDecision {
selected: false,
reason: None,
};
assert_eq!(unselected.reason_text(), "not selected");
}
#[test]
fn effective_target_count_rules() {
assert_eq!(effective_target_count(0, None), 0);
assert_eq!(effective_target_count(0, Some(2)), 0);
assert_eq!(effective_target_count(5, None), 4);
assert_eq!(effective_target_count(20, None), 16);
assert_eq!(effective_target_count(1, Some(2)), 2);
assert_eq!(effective_target_count(10, Some(1)), 1);
assert_eq!(effective_target_count(10, Some(0)), 8);
}
}