klieo_core/response.rs
1//! Typed structured-output trait + parser.
2//!
3//! [`KlieoResponse`] lets callers map an LLM reply onto a strongly-typed
4//! Rust struct. The caller derives the trait (via
5//! `#[derive(KlieoResponse)]` from `klieo-macros`) on a type that also
6//! implements [`serde::de::DeserializeOwned`] and `schemars::JsonSchema`,
7//! then:
8//!
9//! 1. Calls `T::json_schema()` and passes it as
10//! `ResponseFormat::StructuredOutput { schema }` (or
11//! `ResponseFormat::Json { schema }`) on the [`crate::ChatRequest`].
12//! 2. After the LLM responds, runs [`parse_structured`] on the reply
13//! text to recover a typed `T`.
14//!
15//! The parser tolerates the common failure mode where providers wrap
16//! JSON in markdown code fences despite being asked for raw JSON.
17
18use crate::error::Error;
19
20/// Maximum byte length of input accepted by [`parse_structured`].
21///
22/// Multi-megabyte malicious or runaway LLM replies would otherwise be
23/// allocated end-to-end before `serde_json` rejects them. Capping
24/// at 1 MiB bounds memory pressure for compliance-grade deployments
25/// where parser inputs may originate from third-party providers.
26pub const MAX_PARSE_INPUT_BYTES: usize = 1024 * 1024;
27
28/// Trait implemented by types that can be parsed from a structured LLM
29/// response.
30///
31/// Typically derived via `#[derive(KlieoResponse)]` from
32/// `klieo-macros`. Callers must also derive `serde::Deserialize` and
33/// `schemars::JsonSchema` on the same type — the derive only emits the
34/// `KlieoResponse` impl itself and forwards schema generation to
35/// `schemars`.
36///
37/// # Method-name collision with `schemars::JsonSchema`
38///
39/// Both this trait and `schemars::JsonSchema` define a `json_schema`
40/// associated function. When a type derives both, call sites must
41/// disambiguate via fully-qualified syntax:
42///
43/// ```ignore
44/// let schema = <MyType as klieo_core::KlieoResponse>::json_schema();
45/// ```
46///
47/// The collision is intentional: keeping `json_schema()` as the public
48/// surface here matches the obvious naming, and disambiguation is a
49/// one-line cost paid only at the rare site that touches the schema
50/// directly (most callers hand the value off to `ChatRequest`).
51///
52/// # Example (manual impl, no macro)
53///
54/// ```
55/// use klieo_core::response::{KlieoResponse, parse_structured};
56/// use serde::Deserialize;
57/// use serde_json::json;
58///
59/// #[derive(Deserialize)]
60/// struct Greeting {
61/// greeting: String,
62/// }
63///
64/// impl KlieoResponse for Greeting {
65/// fn json_schema() -> serde_json::Value {
66/// json!({
67/// "type": "object",
68/// "properties": { "greeting": { "type": "string" } },
69/// "required": ["greeting"],
70/// })
71/// }
72/// }
73///
74/// let g: Greeting = parse_structured(r#"{"greeting": "hi"}"#).unwrap();
75/// assert_eq!(g.greeting, "hi");
76/// ```
77pub trait KlieoResponse: serde::de::DeserializeOwned + Sized {
78 /// JSON schema describing the shape this type expects.
79 ///
80 /// Pass into `ChatRequest.response_format` as
81 /// `ResponseFormat::StructuredOutput { schema }` (or
82 /// `Json { schema }`).
83 fn json_schema() -> serde_json::Value;
84}
85
86/// Parse an LLM reply (raw text content) into a typed `T`.
87///
88/// Tolerates surrounding markdown code-fences (```` ```json ... ``` ````
89/// or bare ```` ``` ... ``` ````) and leading / trailing whitespace.
90/// Many providers wrap JSON in fences despite being asked for raw JSON;
91/// this helper papers over that.
92///
93/// # Size cap
94///
95/// Inputs larger than [`MAX_PARSE_INPUT_BYTES`] are rejected before any
96/// allocation-heavy parsing runs, so a hostile or runaway provider
97/// cannot push the host into multi-megabyte allocations on every reply.
98///
99/// # Errors
100///
101/// Returns [`Error::BadResponse`] on any deserialization failure. The
102/// payload is intentionally **position-only** (`line` / `col`) — the
103/// raw `serde_json::Error::Display` body would echo attacker- or
104/// LLM-controlled byte fragments into the run-log / OTLP traces, which
105/// matters in compliance deployments. The variant is permanent
106/// (`retryable() == false`); retrying the same reply will fail
107/// identically. Callers wishing to retry should request a fresh
108/// completion with a tighter prompt or schema.
109pub fn parse_structured<T: KlieoResponse>(content: &str) -> Result<T, Error> {
110 if content.len() > MAX_PARSE_INPUT_BYTES {
111 return Err(Error::BadResponse(format!(
112 "structured-output parse error: input exceeds {MAX_PARSE_INPUT_BYTES}-byte cap (got {} bytes)",
113 content.len()
114 )));
115 }
116 let stripped = strip_code_fences(content.trim());
117 serde_json::from_str::<T>(stripped).map_err(|e| {
118 Error::BadResponse(format!(
119 "structured-output parse error at line {} col {}",
120 e.line(),
121 e.column()
122 ))
123 })
124}
125
126/// Strip a leading ```` ```json\n ```` (or ```` ```\n ````) and trailing
127/// ```` ``` ```` if both are present. Returns the original string
128/// otherwise.
129fn strip_code_fences(s: &str) -> &str {
130 let s = s.trim();
131 let Some(rest) = s.strip_prefix("```") else {
132 return s;
133 };
134 // Skip an optional language tag on the opening fence (e.g. `json`).
135 // The fence body starts after the next newline. If there's no
136 // newline, treat the input as un-fenced.
137 let after_lang = match rest.find('\n') {
138 Some(idx) => &rest[idx + 1..],
139 None => return s,
140 };
141 let Some(body) = after_lang.strip_suffix("```") else {
142 return s;
143 };
144 body.trim_end_matches('\n').trim_end()
145}
146
147#[cfg(test)]
148mod tests {
149 use super::*;
150 use serde::Deserialize;
151 use serde_json::json;
152
153 #[derive(Debug, Deserialize, PartialEq)]
154 struct Sample {
155 name: String,
156 }
157
158 impl KlieoResponse for Sample {
159 fn json_schema() -> serde_json::Value {
160 json!({
161 "type": "object",
162 "properties": { "name": { "type": "string" } },
163 "required": ["name"],
164 })
165 }
166 }
167
168 #[test]
169 fn strip_no_fences_passthrough() {
170 assert_eq!(strip_code_fences(r#"{"a":1}"#), r#"{"a":1}"#);
171 }
172
173 #[test]
174 fn strip_json_tagged_fence() {
175 let input = "```json\n{\"a\":1}\n```";
176 assert_eq!(strip_code_fences(input), "{\"a\":1}");
177 }
178
179 #[test]
180 fn strip_bare_fence() {
181 let input = "```\n{\"a\":1}\n```";
182 assert_eq!(strip_code_fences(input), "{\"a\":1}");
183 }
184
185 #[test]
186 fn strip_fence_without_newline_passthrough() {
187 // No newline after opening fence — input is malformed for our
188 // expectations; pass through and let serde fail loudly.
189 let input = "```{\"a\":1}```";
190 assert_eq!(strip_code_fences(input), input);
191 }
192
193 #[test]
194 fn parse_plain_json() {
195 let s: Sample = parse_structured(r#"{"name": "foo"}"#).unwrap();
196 assert_eq!(s, Sample { name: "foo".into() });
197 }
198
199 #[test]
200 fn parse_strips_fences() {
201 let s: Sample = parse_structured("```json\n{\"name\": \"foo\"}\n```").unwrap();
202 assert_eq!(s, Sample { name: "foo".into() });
203 }
204
205 #[test]
206 fn parse_rejects_malformed() {
207 let err = parse_structured::<Sample>("not json").unwrap_err();
208 assert!(matches!(err, Error::BadResponse(_)));
209 assert!(!err.retryable());
210 }
211}