o402 0.1.4

//! Buffer the request body and store [`RequestContext`] in extensions.

use axum::body::Body;
use axum::extract::{Request, State};
use axum::middleware::Next;
use axum::response::Response;
use http_body_util::BodyExt as _;

use super::bill::{self, Bill};
use super::openai_error;
use crate::state::AppState;

/// Parsed fields from a billed request.
#[derive(Clone, Debug, Eq, PartialEq)]
pub(crate) struct RequestContext {
    /// Billing class.
    pub(crate) bill: Bill,
    /// Client-facing model id. `"*"` when the body has none.
    pub(crate) model: String,
    /// `stream`, default false.
    pub(crate) stream: bool,
    /// Optional `max_tokens`.
    pub(crate) max_tokens: Option<u32>,
    /// Optional `max_completion_tokens`.
    pub(crate) max_completion_tokens: Option<u32>,
    /// Optional `max_output_tokens`.
    pub(crate) max_output_tokens: Option<u32>,
}

/// JSON parse failures for billed request bodies.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum ParseError {
    /// Body is not a JSON object.
    InvalidJson,
    /// `model` is missing or empty.
    MissingModel,
    /// `model` is present but not a string.
    ModelNotString,
    /// `stream` is present but not a boolean.
    StreamNotBool,
}

/// Buffer JSON when needed, restore the body, store [`RequestContext`].
pub(super) async fn extract_context(
    State(_state): State<AppState>,
    request: Request,
    next: Next,
) -> Response {
    let path = request.uri().path();
    if path == "/health" || path == "/ready" || path == "/v1/pricing" {
        return next.run(request).await;
    }
    let bill = bill::classify(request.method(), path, request.headers());
    match bill {
        Bill::Reject => {
            return openai_error::not_found(format!("Invalid URL ({})", request.uri().path()));
        }
        Bill::Unpaid => return next.run(request).await,
        Bill::Exact | Bill::Upto => {}
    }
    let (parts, body) = request.into_parts();
    let bytes = match body.collect().await {
        Ok(collected) => collected.to_bytes(),
        Err(_) => return openai_error::invalid_request("failed to read request body"),
    };
    let ctx = match parse_request_context(bill, &bytes) {
        Ok(ctx) => ctx,
        Err(error) => return parse_error_response(error),
    };
    let mut restored = Request::from_parts(parts, Body::from(bytes));
    restored.extensions_mut().insert(ctx);
    next.run(restored).await
}

fn parse_request_context(bill: Bill, bytes: &[u8]) -> Result<RequestContext, ParseError> {
    if bytes.is_empty() {
        return match bill {
            Bill::Upto => Err(ParseError::InvalidJson),
            Bill::Exact => Ok(RequestContext {
                bill,
                model: "*".to_owned(),
                stream: false,
                max_tokens: None,
                max_completion_tokens: None,
                max_output_tokens: None,
            }),
            Bill::Unpaid | Bill::Reject => unreachable!("unpaid/reject skip parse"),
        };
    }
    let value: serde_json::Value = match serde_json::from_slice(bytes) {
        Ok(value) => value,
        Err(_) if bill == Bill::Exact => {
            return Ok(RequestContext {
                bill,
                model: "*".to_owned(),
                stream: false,
                max_tokens: None,
                max_completion_tokens: None,
                max_output_tokens: None,
            });
        }
        Err(_) => return Err(ParseError::InvalidJson),
    };
    if !value.is_object() {
        return match bill {
            Bill::Exact => Ok(RequestContext {
                bill,
                model: "*".to_owned(),
                stream: false,
                max_tokens: None,
                max_completion_tokens: None,
                max_output_tokens: None,
            }),
            _ => Err(ParseError::InvalidJson),
        };
    }
    let model = match value.get("model") {
        None | Some(serde_json::Value::Null) if bill == Bill::Exact => "*".to_owned(),
        None => return Err(ParseError::MissingModel),
        Some(serde_json::Value::String(model)) if model.is_empty() => {
            return Err(ParseError::MissingModel);
        }
        Some(serde_json::Value::String(model)) => model.clone(),
        Some(_) => return Err(ParseError::ModelNotString),
    };
    let stream = match value.get("stream") {
        None | Some(serde_json::Value::Null) => false,
        Some(serde_json::Value::Bool(flag)) => *flag,
        Some(_) => return Err(ParseError::StreamNotBool),
    };
    Ok(RequestContext {
        bill,
        model,
        stream,
        max_tokens: json_u32(&value, "max_tokens"),
        max_completion_tokens: json_u32(&value, "max_completion_tokens"),
        max_output_tokens: json_u32(&value, "max_output_tokens"),
    })
}

fn json_u32(value: &serde_json::Value, key: &str) -> Option<u32> {
    value
        .get(key)
        .and_then(serde_json::Value::as_u64)
        .and_then(|n| u32::try_from(n).ok())
}

fn parse_error_response(error: ParseError) -> Response {
    match error {
        ParseError::InvalidJson => {
            openai_error::invalid_request("We could not parse the JSON body of your request")
        }
        ParseError::MissingModel => {
            openai_error::invalid_request("you must provide a model parameter")
        }
        ParseError::ModelNotString => openai_error::invalid_request("model must be a string"),
        ParseError::StreamNotBool => openai_error::invalid_request("stream must be a boolean"),
    }
}

#[cfg(test)]
mod tests {
    use super::{Bill, ParseError, parse_request_context};

    #[test]
    fn parses_model_stream_and_token_caps() {
        let ctx = parse_request_context(
            Bill::Upto,
            br#"{"model":"gpt-4o-mini","stream":true,"max_tokens":16,"max_completion_tokens":32,"max_output_tokens":8}"#,
        )
        .expect("parse");
        assert_eq!(ctx.model, "gpt-4o-mini", "model");
        assert!(ctx.stream, "stream");
        assert_eq!(ctx.max_tokens, Some(16), "max_tokens");
        assert_eq!(ctx.max_completion_tokens, Some(32), "max_completion_tokens");
        assert_eq!(ctx.max_output_tokens, Some(8), "max_output_tokens");
    }

    #[test]
    fn stream_defaults_false() {
        let ctx = parse_request_context(Bill::Upto, br#"{"model":"gpt-4o-mini"}"#).expect("parse");
        assert!(!ctx.stream, "default stream");
        assert_eq!(ctx.max_tokens, None, "max_tokens");
    }

    #[test]
    fn missing_model_is_error() {
        assert_eq!(
            parse_request_context(Bill::Upto, br#"{"messages":[]}"#),
            Err(ParseError::MissingModel),
            "missing"
        );
    }

    #[test]
    fn invalid_json_is_error() {
        assert_eq!(
            parse_request_context(Bill::Upto, b"not-json"),
            Err(ParseError::InvalidJson),
            "json"
        );
    }

    #[test]
    fn stream_must_be_bool() {
        assert_eq!(
            parse_request_context(Bill::Upto, br#"{"model":"gpt-4o-mini","stream":1}"#),
            Err(ParseError::StreamNotBool),
            "stream"
        );
    }
}