o402 0.1.4

//! `GET /v1/pricing` from config rates.

use axum::Json;
use axum::extract::State;
use serde::Serialize;

use crate::config::{ModelConfig, PricingDefault};
use crate::state::AppState;

/// Pricing catalog body.
#[derive(Serialize)]
pub(super) struct PricingBody {
    /// One row per `[[models]]` entry.
    models: Vec<PricingModel>,
}

/// Config rates for one model. Missing fields are omitted.
#[derive(Serialize)]
struct PricingModel {
    /// Client-facing id.
    id: String,
    /// `owned_by` from config, else `system`.
    owned_by: String,
    /// `exact` or `upto`.
    #[serde(skip_serializing_if = "Option::is_none")]
    scheme: Option<String>,
    /// Exact per-request price.
    #[serde(skip_serializing_if = "Option::is_none")]
    price: Option<String>,
    /// Request floor.
    #[serde(skip_serializing_if = "Option::is_none")]
    request_floor: Option<String>,
    /// Uncached input rate.
    #[serde(skip_serializing_if = "Option::is_none")]
    input_per_million: Option<String>,
    /// Non-reasoning output rate.
    #[serde(skip_serializing_if = "Option::is_none")]
    output_per_million: Option<String>,
    /// Cached input rate.
    #[serde(skip_serializing_if = "Option::is_none")]
    cached_input_per_million: Option<String>,
    /// Reasoning output rate.
    #[serde(skip_serializing_if = "Option::is_none")]
    reasoning_per_million: Option<String>,
    /// Ceiling multiplier.
    #[serde(skip_serializing_if = "Option::is_none")]
    ceiling_multiplier: Option<String>,
    /// Hard ceiling.
    #[serde(skip_serializing_if = "Option::is_none")]
    max_ceiling: Option<String>,
    /// Max input tokens.
    #[serde(skip_serializing_if = "Option::is_none")]
    max_input_tokens: Option<u32>,
    /// Default max output tokens.
    #[serde(skip_serializing_if = "Option::is_none")]
    default_max_output_tokens: Option<u32>,
}

/// Return configured rates. Empty when there is no catalog.
#[allow(clippy::unused_async, reason = "axum handler")]
pub(super) async fn list(State(state): State<AppState>) -> Json<PricingBody> {
    let default = state.pricing().and_then(|pricing| pricing.default.as_ref());
    Json(PricingBody {
        models: state
            .models()
            .iter()
            .map(|model| pricing_model(model, default))
            .collect(),
    })
}

fn pricing_model(model: &ModelConfig, default: Option<&PricingDefault>) -> PricingModel {
    PricingModel {
        id: model.id.clone(),
        owned_by: model.owned_by.as_deref().unwrap_or("system").to_owned(),
        scheme: inherit_str(model.scheme.as_ref(), default, |row| row.scheme.as_ref()),
        price: inherit_str(model.price.as_ref(), default, |row| row.price.as_ref()),
        request_floor: inherit_str(model.request_floor.as_ref(), default, |row| {
            row.request_floor.as_ref()
        }),
        input_per_million: inherit_str(model.input_per_million.as_ref(), default, |row| {
            row.input_per_million.as_ref()
        }),
        output_per_million: inherit_str(model.output_per_million.as_ref(), default, |row| {
            row.output_per_million.as_ref()
        }),
        cached_input_per_million: inherit_str(
            model.cached_input_per_million.as_ref(),
            default,
            |row| row.cached_input_per_million.as_ref(),
        ),
        reasoning_per_million: inherit_str(model.reasoning_per_million.as_ref(), default, |row| {
            row.reasoning_per_million.as_ref()
        }),
        ceiling_multiplier: inherit_str(model.ceiling_multiplier.as_ref(), default, |row| {
            row.ceiling_multiplier.as_ref()
        }),
        max_ceiling: inherit_str(model.max_ceiling.as_ref(), default, |row| {
            row.max_ceiling.as_ref()
        }),
        max_input_tokens: inherit_u32(model.max_input_tokens, default, |row| row.max_input_tokens),
        default_max_output_tokens: inherit_u32(model.default_max_output_tokens, default, |row| {
            row.default_max_output_tokens
        }),
    }
}

fn inherit_str(
    model: Option<&String>,
    default: Option<&PricingDefault>,
    pick: impl Fn(&PricingDefault) -> Option<&String>,
) -> Option<String> {
    model.or_else(|| default.and_then(pick)).cloned()
}

fn inherit_u32(
    model: Option<u32>,
    default: Option<&PricingDefault>,
    pick: impl Fn(&PricingDefault) -> Option<u32>,
) -> Option<u32> {
    model.or_else(|| default.and_then(pick))
}

#[cfg(test)]
mod tests {
    use axum::body::Body;
    use axum::http::{Request, StatusCode};
    use tower::ServiceExt;

    use crate::config::Config;
    use crate::http::app;

    #[tokio::test]
    async fn inherits_pricing_default() {
        let toml = r#"
[payment]
enabled = false

[pricing.default]
scheme = "exact"
price = "0.001"

[[upstreams]]
name = "stub"
base_url = "http://127.0.0.1:9"
api_key = "sk-test"

[[models]]
id = "gpt-4o-mini"
upstream = "stub"
owned_by = "openai"
"#;
        let response = app(Config::from_toml_str(toml).expect("config"))
            .expect("router")
            .oneshot(
                Request::get("/v1/pricing")
                    .body(Body::empty())
                    .expect("request"),
            )
            .await
            .expect("response");
        assert_eq!(response.status(), StatusCode::OK, "status");
        let bytes = axum::body::to_bytes(response.into_body(), 4096)
            .await
            .expect("body");
        let body: serde_json::Value = serde_json::from_slice(&bytes).expect("json");
        let models = body
            .get("models")
            .and_then(serde_json::Value::as_array)
            .expect("models");
        assert_eq!(models.len(), 1, "len");
        let model = models.first().expect("model");
        assert_eq!(
            model.get("id").and_then(serde_json::Value::as_str),
            Some("gpt-4o-mini"),
            "id"
        );
        assert_eq!(
            model.get("scheme").and_then(serde_json::Value::as_str),
            Some("exact"),
            "scheme"
        );
        assert_eq!(
            model.get("price").and_then(serde_json::Value::as_str),
            Some("0.001"),
            "price"
        );
    }

    #[tokio::test]
    async fn shows_upto_rates() {
        let toml = r#"
[payment]
enabled = false

[pricing.default]
scheme = "upto"
request_floor = "0.00001"
input_per_million = "0.15"
output_per_million = "0.60"
cached_input_per_million = "0.075"
reasoning_per_million = "0.60"
max_ceiling = "5.00"
max_input_tokens = 128000
default_max_output_tokens = 16384

[[upstreams]]
name = "stub"
base_url = "http://127.0.0.1:9"
api_key = "sk-test"

[[models]]
id = "gpt-4o-mini"
upstream = "stub"
owned_by = "openai"
"#;
        let response = app(Config::from_toml_str(toml).expect("config"))
            .expect("router")
            .oneshot(
                Request::get("/v1/pricing")
                    .body(Body::empty())
                    .expect("request"),
            )
            .await
            .expect("response");
        assert_eq!(response.status(), StatusCode::OK, "status");
        let bytes = axum::body::to_bytes(response.into_body(), 4096)
            .await
            .expect("body");
        let body: serde_json::Value = serde_json::from_slice(&bytes).expect("json");
        let model = body
            .get("models")
            .and_then(serde_json::Value::as_array)
            .and_then(|models| models.first())
            .expect("model");
        assert_eq!(
            model.get("scheme").and_then(serde_json::Value::as_str),
            Some("upto"),
            "scheme"
        );
        assert_eq!(
            model
                .get("input_per_million")
                .and_then(serde_json::Value::as_str),
            Some("0.15"),
            "input"
        );
        assert_eq!(
            model
                .get("output_per_million")
                .and_then(serde_json::Value::as_str),
            Some("0.60"),
            "output"
        );
        assert_eq!(
            model
                .get("cached_input_per_million")
                .and_then(serde_json::Value::as_str),
            Some("0.075"),
            "cached"
        );
    }

    #[tokio::test]
    async fn embeddings_explicit_zero_output_rate() {
        let toml = r#"
[payment]
enabled = false

[pricing.default]
scheme = "upto"
input_per_million = "0.15"
output_per_million = "0.60"
reasoning_per_million = "0.60"
default_max_output_tokens = 16384

[[upstreams]]
name = "stub"
base_url = "http://127.0.0.1:9"
api_key = "sk-test"

[[models]]
id = "text-embedding-3-small"
upstream = "stub"
scheme = "upto"
input_per_million = "0.02"
output_per_million = "0"
max_input_tokens = 8191
"#;
        let response = app(Config::from_toml_str(toml).expect("config"))
            .expect("router")
            .oneshot(
                Request::get("/v1/pricing")
                    .body(Body::empty())
                    .expect("request"),
            )
            .await
            .expect("response");
        assert_eq!(response.status(), StatusCode::OK, "status");
        let bytes = axum::body::to_bytes(response.into_body(), 4096)
            .await
            .expect("body");
        let body: serde_json::Value = serde_json::from_slice(&bytes).expect("json");
        let model = body
            .get("models")
            .and_then(serde_json::Value::as_array)
            .and_then(|models| models.first())
            .expect("model");
        assert_eq!(
            model
                .get("input_per_million")
                .and_then(serde_json::Value::as_str),
            Some("0.02"),
            "input"
        );
        assert_eq!(
            model
                .get("output_per_million")
                .and_then(serde_json::Value::as_str),
            Some("0"),
            "explicit zero output"
        );
    }

    #[tokio::test]
    async fn empty_when_no_models() {
        let toml = r#"
[payment]
enabled = false

[[upstreams]]
name = "stub"
base_url = "http://127.0.0.1:9"
api_key = "sk-test"
"#;
        let response = app(Config::from_toml_str(toml).expect("config"))
            .expect("router")
            .oneshot(
                Request::get("/v1/pricing")
                    .body(Body::empty())
                    .expect("request"),
            )
            .await
            .expect("response");
        assert_eq!(response.status(), StatusCode::OK, "status");
        let bytes = axum::body::to_bytes(response.into_body(), 4096)
            .await
            .expect("body");
        let body: serde_json::Value = serde_json::from_slice(&bytes).expect("json");
        assert_eq!(
            body.get("models")
                .and_then(serde_json::Value::as_array)
                .map(Vec::len),
            Some(0),
            "empty"
        );
    }
}