kcode-k1-chat-websearch-request 0.1.0

use kcode_k1_codex_websearch::Request;
use serde::{Deserialize, Deserializer};
use std::time::Instant;

const INVALID_ARGUMENTS: &str = "WebSearch failed: invalid arguments";
const UNSUPPORTED_ROUTE: &str =
    "WebSearch failed: unsupported model route; expected model to begin with \"codex/\"";

#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct PublicWebSearch {
    query: String,
    model: String,
    #[serde(default, deserialize_with = "present_string")]
    reasoning_effort: Option<String>,
}

pub fn parse(arguments: &str, deadline: Instant) -> Result<Request, String> {
    let public: PublicWebSearch =
        serde_json::from_str(arguments).map_err(|_| INVALID_ARGUMENTS.to_owned())?;
    let model = public
        .model
        .strip_prefix("codex/")
        .ok_or_else(|| UNSUPPORTED_ROUTE.to_owned())?;
    Ok(Request {
        query: public.query,
        model: model.to_owned(),
        reasoning_effort: public
            .reasoning_effort
            .unwrap_or_else(|| "medium".to_owned()),
        deadline,
    })
}

fn present_string<'de, D: Deserializer<'de>>(deserializer: D) -> Result<Option<String>, D::Error> {
    String::deserialize(deserializer).map(Some)
}

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

    #[test]
    fn valid_values_are_preserved_and_effort_defaults() {
        let deadline = Instant::now() + Duration::from_secs(10);
        let request = parse(
            r#"{"query":"  exact\nquery  ","model":"codex/model","reasoning_effort":"high"}"#,
            deadline,
        )
        .unwrap();
        assert_eq!(request.query, "  exact\nquery  ");
        assert_eq!(request.model, "model");
        assert_eq!(request.reasoning_effort, "high");
        assert_eq!(request.deadline, deadline);

        let request = parse(r#"{"query":"","model":"codex/m"}"#, deadline).unwrap();
        assert_eq!(request.query, "");
        assert_eq!(request.reasoning_effort, "medium");
    }

    #[test]
    fn invalid_shapes_are_rejected_uniformly() {
        for arguments in [
            "null",
            "[]",
            "{}",
            r#"{"query":"q"}"#,
            r#"{"model":"codex/m"}"#,
            r#"{"query":null,"model":"codex/m"}"#,
            r#"{"query":"q","model":7}"#,
            r#"{"query":"q","model":"codex/m","reasoning_effort":null}"#,
            r#"{"query":"q","model":"codex/m","reasoning_effort":7}"#,
            r#"{"query":"q","model":"codex/m","extra":true}"#,
            r#"{"question":"q","model":"codex/m"}"#,
        ] {
            let error = match parse(arguments, Instant::now()) {
                Err(error) => error,
                Ok(_) => panic!("invalid WebSearch arguments were accepted"),
            };
            assert_eq!(error, INVALID_ARGUMENTS);
        }
    }

    #[test]
    fn model_route_is_exact_case_sensitive_and_stripped_once() {
        for model in ["model", "Codex/model", "codexx/model"] {
            let arguments = format!(r#"{{"query":"q","model":"{model}"}}"#);
            let error = match parse(&arguments, Instant::now()) {
                Err(error) => error,
                Ok(_) => panic!("unsupported WebSearch route was accepted"),
            };
            assert_eq!(error, UNSUPPORTED_ROUTE);
        }

        let request = parse(
            r#"{"query":"q","model":"codex/codex/model"}"#,
            Instant::now(),
        )
        .unwrap();
        assert_eq!(request.model, "codex/model");
    }
}