klieo-core 0.6.0

//! Typed structured-output trait + parser.
//!
//! [`KlieoResponse`] lets callers map an LLM reply onto a strongly-typed
//! Rust struct. The caller derives the trait (via
//! `#[derive(KlieoResponse)]` from `klieo-macros`) on a type that also
//! implements [`serde::de::DeserializeOwned`] and `schemars::JsonSchema`,
//! then:
//!
//! 1. Calls `T::json_schema()` and passes it as
//!    `ResponseFormat::StructuredOutput { schema }` (or
//!    `ResponseFormat::Json { schema }`) on the [`crate::ChatRequest`].
//! 2. After the LLM responds, runs [`parse_structured`] on the reply
//!    text to recover a typed `T`.
//!
//! The parser tolerates the common failure mode where providers wrap
//! JSON in markdown code fences despite being asked for raw JSON.

use crate::error::Error;

/// Maximum byte length of input accepted by [`parse_structured`].
///
/// Multi-megabyte malicious or runaway LLM replies would otherwise be
/// allocated end-to-end before `serde_json` rejects them. Capping
/// at 1 MiB bounds memory pressure for compliance-grade deployments
/// where parser inputs may originate from third-party providers.
pub const MAX_PARSE_INPUT_BYTES: usize = 1024 * 1024;

/// Trait implemented by types that can be parsed from a structured LLM
/// response.
///
/// Typically derived via `#[derive(KlieoResponse)]` from
/// `klieo-macros`. Callers must also derive `serde::Deserialize` and
/// `schemars::JsonSchema` on the same type — the derive only emits the
/// `KlieoResponse` impl itself and forwards schema generation to
/// `schemars`.
///
/// # Method-name collision with `schemars::JsonSchema`
///
/// Both this trait and `schemars::JsonSchema` define a `json_schema`
/// associated function. When a type derives both, call sites must
/// disambiguate via fully-qualified syntax:
///
/// ```ignore
/// let schema = <MyType as klieo_core::KlieoResponse>::json_schema();
/// ```
///
/// The collision is intentional: keeping `json_schema()` as the public
/// surface here matches the obvious naming, and disambiguation is a
/// one-line cost paid only at the rare site that touches the schema
/// directly (most callers hand the value off to `ChatRequest`).
///
/// # Example (manual impl, no macro)
///
/// ```
/// use klieo_core::response::{KlieoResponse, parse_structured};
/// use serde::Deserialize;
/// use serde_json::json;
///
/// #[derive(Deserialize)]
/// struct Greeting {
///     greeting: String,
/// }
///
/// impl KlieoResponse for Greeting {
///     fn json_schema() -> serde_json::Value {
///         json!({
///             "type": "object",
///             "properties": { "greeting": { "type": "string" } },
///             "required": ["greeting"],
///         })
///     }
/// }
///
/// let g: Greeting = parse_structured(r#"{"greeting": "hi"}"#).unwrap();
/// assert_eq!(g.greeting, "hi");
/// ```
pub trait KlieoResponse: serde::de::DeserializeOwned + Sized {
    /// JSON schema describing the shape this type expects.
    ///
    /// Pass into `ChatRequest.response_format` as
    /// `ResponseFormat::StructuredOutput { schema }` (or
    /// `Json { schema }`).
    fn json_schema() -> serde_json::Value;
}

/// Parse an LLM reply (raw text content) into a typed `T`.
///
/// Tolerates surrounding markdown code-fences (```` ```json ... ``` ````
/// or bare ```` ``` ... ``` ````) and leading / trailing whitespace.
/// Many providers wrap JSON in fences despite being asked for raw JSON;
/// this helper papers over that.
///
/// # Size cap
///
/// Inputs larger than [`MAX_PARSE_INPUT_BYTES`] are rejected before any
/// allocation-heavy parsing runs, so a hostile or runaway provider
/// cannot push the host into multi-megabyte allocations on every reply.
///
/// # Errors
///
/// Returns [`Error::BadResponse`] on any deserialization failure. The
/// payload is intentionally **position-only** (`line` / `col`) — the
/// raw `serde_json::Error::Display` body would echo attacker- or
/// LLM-controlled byte fragments into the run-log / OTLP traces, which
/// matters in compliance deployments. The variant is permanent
/// (`retryable() == false`); retrying the same reply will fail
/// identically. Callers wishing to retry should request a fresh
/// completion with a tighter prompt or schema.
pub fn parse_structured<T: KlieoResponse>(content: &str) -> Result<T, Error> {
    if content.len() > MAX_PARSE_INPUT_BYTES {
        return Err(Error::BadResponse(format!(
            "structured-output parse error: input exceeds {MAX_PARSE_INPUT_BYTES}-byte cap (got {} bytes)",
            content.len()
        )));
    }
    let stripped = strip_code_fences(content.trim());
    serde_json::from_str::<T>(stripped).map_err(|e| {
        Error::BadResponse(format!(
            "structured-output parse error at line {} col {}",
            e.line(),
            e.column()
        ))
    })
}

/// Strip a leading ```` ```json\n ```` (or ```` ```\n ````) and trailing
/// ```` ``` ```` if both are present. Returns the original string
/// otherwise.
fn strip_code_fences(s: &str) -> &str {
    let s = s.trim();
    let Some(rest) = s.strip_prefix("```") else {
        return s;
    };
    // Skip an optional language tag on the opening fence (e.g. `json`).
    // The fence body starts after the next newline. If there's no
    // newline, treat the input as un-fenced.
    let after_lang = match rest.find('\n') {
        Some(idx) => &rest[idx + 1..],
        None => return s,
    };
    let Some(body) = after_lang.strip_suffix("```") else {
        return s;
    };
    body.trim_end_matches('\n').trim_end()
}

#[cfg(test)]
mod tests {
    use super::*;
    use serde::Deserialize;
    use serde_json::json;

    #[derive(Debug, Deserialize, PartialEq)]
    struct Sample {
        name: String,
    }

    impl KlieoResponse for Sample {
        fn json_schema() -> serde_json::Value {
            json!({
                "type": "object",
                "properties": { "name": { "type": "string" } },
                "required": ["name"],
            })
        }
    }

    #[test]
    fn strip_no_fences_passthrough() {
        assert_eq!(strip_code_fences(r#"{"a":1}"#), r#"{"a":1}"#);
    }

    #[test]
    fn strip_json_tagged_fence() {
        let input = "```json\n{\"a\":1}\n```";
        assert_eq!(strip_code_fences(input), "{\"a\":1}");
    }

    #[test]
    fn strip_bare_fence() {
        let input = "```\n{\"a\":1}\n```";
        assert_eq!(strip_code_fences(input), "{\"a\":1}");
    }

    #[test]
    fn strip_fence_without_newline_passthrough() {
        // No newline after opening fence — input is malformed for our
        // expectations; pass through and let serde fail loudly.
        let input = "```{\"a\":1}```";
        assert_eq!(strip_code_fences(input), input);
    }

    #[test]
    fn parse_plain_json() {
        let s: Sample = parse_structured(r#"{"name": "foo"}"#).unwrap();
        assert_eq!(s, Sample { name: "foo".into() });
    }

    #[test]
    fn parse_strips_fences() {
        let s: Sample = parse_structured("```json\n{\"name\": \"foo\"}\n```").unwrap();
        assert_eq!(s, Sample { name: "foo".into() });
    }

    #[test]
    fn parse_rejects_malformed() {
        let err = parse_structured::<Sample>("not json").unwrap_err();
        assert!(matches!(err, Error::BadResponse(_)));
        assert!(!err.retryable());
    }
}