garudust-core 0.3.0

/// Per-million-token prices (USD) for known models.
/// Prices are approximate and updated manually. Models not listed here return `None`.
struct ModelPrice {
    input_per_m: f64,
    output_per_m: f64,
}

fn lookup(model: &str) -> Option<ModelPrice> {
    // Normalise: strip provider prefix ("anthropic/claude-...") and minor version suffixes.
    let name = model
        .rsplit_once('/')
        .map_or(model, |(_, n)| n)
        .to_lowercase();

    // Match by prefix so minor variants (e.g. -20251001) are covered.
    let p = if name.starts_with("claude-opus-4") {
        ModelPrice {
            input_per_m: 15.0,
            output_per_m: 75.0,
        }
    } else if name.starts_with("claude-sonnet-4") {
        ModelPrice {
            input_per_m: 3.0,
            output_per_m: 15.0,
        }
    } else if name.starts_with("claude-haiku-4") {
        ModelPrice {
            input_per_m: 0.8,
            output_per_m: 4.0,
        }
    } else if name.starts_with("claude-opus-3") {
        ModelPrice {
            input_per_m: 15.0,
            output_per_m: 75.0,
        }
    } else if name.starts_with("claude-sonnet-3") {
        ModelPrice {
            input_per_m: 3.0,
            output_per_m: 15.0,
        }
    } else if name.starts_with("claude-haiku-3") {
        ModelPrice {
            input_per_m: 0.25,
            output_per_m: 1.25,
        }
    } else if name.starts_with("gpt-4o-mini") {
        ModelPrice {
            input_per_m: 0.15,
            output_per_m: 0.60,
        }
    } else if name.starts_with("gpt-4o") {
        ModelPrice {
            input_per_m: 2.50,
            output_per_m: 10.0,
        }
    } else if name.starts_with("gpt-4-turbo") || name.starts_with("gpt-4-1106") {
        ModelPrice {
            input_per_m: 10.0,
            output_per_m: 30.0,
        }
    } else if name.starts_with("gpt-3.5") {
        ModelPrice {
            input_per_m: 0.50,
            output_per_m: 1.50,
        }
    } else if name.starts_with("gemini-2.5-pro") {
        ModelPrice {
            input_per_m: 1.25,
            output_per_m: 10.0,
        }
    } else if name.starts_with("gemini-2.5-flash") {
        ModelPrice {
            input_per_m: 0.15,
            output_per_m: 0.60,
        }
    } else if name.starts_with("gemini-2.0-flash") {
        ModelPrice {
            input_per_m: 0.10,
            output_per_m: 0.40,
        }
    } else if name.starts_with("gemini-1.5-pro") {
        ModelPrice {
            input_per_m: 3.50,
            output_per_m: 10.50,
        }
    } else if name.starts_with("gemini-1.5-flash") {
        ModelPrice {
            input_per_m: 0.075,
            output_per_m: 0.30,
        }
    } else {
        return None;
    };
    Some(p)
}

/// Estimate USD cost for a completed task.
/// Returns `None` if the model is not in the price table.
pub fn estimate_cost_usd(model: &str, input_tokens: u32, output_tokens: u32) -> Option<f64> {
    let p = lookup(model)?;
    let cost = (f64::from(input_tokens) / 1_000_000.0) * p.input_per_m
        + (f64::from(output_tokens) / 1_000_000.0) * p.output_per_m;
    Some(cost)
}

/// Format a usage footer line: `[N iter | Xin Yout tok | ~$Z @ model]`
pub fn usage_footer(model: &str, iters: u32, input_tokens: u32, output_tokens: u32) -> String {
    let total = input_tokens + output_tokens;
    let cost_part = estimate_cost_usd(model, input_tokens, output_tokens)
        .map(|c| format!(" | ~${c:.4}"))
        .unwrap_or_default();
    let short_model = model.rsplit_once('/').map_or(model, |(_, n)| n);
    format!(
        "[{iters} iter | {input_tokens}in {output_tokens}out {total}tok{cost_part} @ {short_model}]"
    )
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn known_model_has_cost() {
        let c = estimate_cost_usd("anthropic/claude-sonnet-4-6", 10_000, 2_000);
        assert!(c.is_some());
        // 10k in @ $3/M + 2k out @ $15/M = $0.030 + $0.030 = $0.060
        let c = c.unwrap();
        assert!((c - 0.060).abs() < 0.001, "expected ~$0.06, got {c}");
    }

    #[test]
    fn unknown_model_returns_none() {
        assert!(estimate_cost_usd("local/mistral-7b", 1000, 500).is_none());
    }

    #[test]
    fn footer_includes_iter_and_tokens() {
        let footer = usage_footer("claude-sonnet-4-6", 3, 1000, 500);
        assert!(footer.contains("3 iter"));
        assert!(footer.contains("1000in"));
        assert!(footer.contains("500out"));
        assert!(footer.contains("1500tok"));
    }

    #[test]
    fn footer_strips_provider_prefix() {
        let footer = usage_footer("anthropic/claude-haiku-4-5", 1, 100, 50);
        assert!(footer.contains("@ claude-haiku-4-5"));
        assert!(!footer.contains("anthropic/"));
    }
}