1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
//! JSON-Schema-constrained output for the [`TextModel`] port (issue #580):
//! ask for a value that conforms to a schema, and get either a conforming
//! value or an error — never a half-parsed one.
//!
//! [`Prompt::json_schema`] is already a *request*; what was missing was the
//! guarantee. [`TextModelExt`] is that guarantee, blanket-implemented so an
//! adapter cannot skip it: it checks the schema before the call, asks for it,
//! validates the answer (native [`Completion::json`] or text parsed as JSON)
//! and allows one repair retry. The validator is the tool loop's narrow,
//! fail-closed subset. See the trait and its methods for the full contract.
use async_trait::async_trait;
use schemars::{JsonSchema, Schema};
use serde::de::DeserializeOwned;
use serde_json::Value;
use crate::ports::{Completion, Prompt, TextModel, TextModelError, Turn};
use crate::surface::schema_for;
use crate::tool_loop::{validate_instance, validate_schema};
/// Structured output on any [`TextModel`] (issue #580). Blanket-implemented
/// for every `TextModel`, so an adapter cannot override — and cannot skip —
/// the schema check, the validation, or the repair retry; import the trait
/// to call the methods.
#[async_trait]
pub trait TextModelExt: TextModel {
/// Completes `prompt` and returns a value that conforms to `schema`.
///
/// The schema is checked on the supported subset before the first call,
/// the prompt's [`Prompt::json_schema`] is set to it, and the answer —
/// native structured output, or text parsed as JSON — is validated
/// against it. On failure the model gets one repair retry; a second
/// failure is [`TextModelError::SchemaViolation`].
///
/// # Errors
///
/// - [`TextModelError::InvalidSchema`] if `schema` uses a keyword the
/// validator does not implement; the model is never called.
/// - [`TextModelError::SchemaViolation`] if the answer still does not
/// conform after the repair retry. The message names the failed
/// path and reason only, never the model's output.
/// - Whatever `complete` itself returned, unchanged: `NotConfigured`,
/// `Rejected`, `Transient`, `Transport`, `Unsupported`.
async fn complete_json(&self, prompt: Prompt, schema: &Schema)
-> Result<Value, TextModelError>;
/// Completes `prompt` and deserializes the conforming value into `T`,
/// using `T`'s own schema ([`schema_for`]) as the constraint — the
/// typed form of [`TextModelExt::complete_json`].
///
/// # Errors
///
/// The same as [`TextModelExt::complete_json`], plus
/// [`TextModelError::SchemaViolation`] if the conforming value does not
/// deserialize into `T` (a shape the schema generator and the
/// deserializer disagree on).
async fn complete_as<T>(&self, prompt: Prompt) -> Result<T, TextModelError>
where
T: JsonSchema + DeserializeOwned + Send;
}
#[async_trait]
impl<M: TextModel + ?Sized> TextModelExt for M {
async fn complete_json(
&self,
mut prompt: Prompt,
schema: &Schema,
) -> Result<Value, TextModelError> {
let schema = schema.as_value().clone();
// Fail closed before the model is asked: an unsupported keyword, or a
// root that is not an object, is the caller's mistake.
validate_schema(&schema).map_err(TextModelError::InvalidSchema)?;
// Both native paths require an object root: Anthropic's forced tool
// takes an object `input_schema`, OpenAI's strict `response_format`
// an object. A non-object root has no provider that could honour it.
if schema.get("type").and_then(Value::as_str) != Some("object") {
return Err(TextModelError::InvalidSchema(
"the schema's root type must be \"object\"".to_owned(),
));
}
prompt.json_schema = Some(schema.clone());
let completion = self.complete(&prompt).await?;
match checked(&schema, &completion) {
Ok(value) => Ok(value),
Err(reason) => {
// One repair retry: show the model what it said and what was
// wrong with it, and ask once more for only the JSON.
prompt.messages.push(Turn::assistant(said(&completion)));
prompt
.messages
.push(Turn::user(repair_instruction(&reason)));
let retry = self.complete(&prompt).await?;
checked(&schema, &retry).map_err(TextModelError::SchemaViolation)
}
}
}
async fn complete_as<T>(&self, prompt: Prompt) -> Result<T, TextModelError>
where
T: JsonSchema + DeserializeOwned + Send,
{
let schema = schema_for::<T>();
let value = self.complete_json(prompt, &schema).await?;
serde_json::from_value(value).map_err(|_| {
// The serde error quotes the offending value, which may be the
// personal data the schema was there to structure; the schema
// already checked the shape, so name the failure without it.
TextModelError::SchemaViolation(
"the validated value did not deserialize into the requested type".to_owned(),
)
})
}
}
/// The candidate value from `completion`, or the reason it is unusable:
/// [`Completion::json`] when the provider honoured the schema request, else
/// [`Completion::text`] parsed as JSON with any markdown fence stripped.
fn checked(schema: &Value, completion: &Completion) -> Result<Value, String> {
let candidate = match &completion.json {
Some(json) => json.clone(),
None => serde_json::from_str(strip_code_fence(&completion.text))
.map_err(|_| "the reply was not valid JSON".to_owned())?,
};
validate_instance(schema, &candidate)?;
Ok(candidate)
}
/// What the model said, for quoting back as the assistant turn of a repair:
/// its text, its parsed JSON when a native provider put nothing in the text,
/// or a placeholder — an empty assistant turn is refused by some providers.
fn said(completion: &Completion) -> String {
if !completion.text.is_empty() {
return completion.text.clone();
}
match &completion.json {
Some(json) => json.to_string(),
None => "(empty reply)".to_owned(),
}
}
/// The user turn a repair asks for: the failure, and the one thing to do
/// about it.
fn repair_instruction(reason: &str) -> String {
format!(
"That reply did not conform to the required JSON schema ({reason}). \
Reply with only a JSON value that conforms to the schema, and nothing else."
)
}
/// Strips one wrapping markdown code fence, tolerating the language tag
/// (`json`), a bare fence, and plain whitespace-only padding — the shape a
/// model without native structured output wraps its JSON in.
fn strip_code_fence(text: &str) -> &str {
let trimmed = text.trim();
let Some(after_open) = trimmed.strip_prefix("```") else {
return trimmed;
};
let body = match after_open.split_once('\n') {
Some((_, body)) => body,
None => after_open,
};
match body.rfind("```") {
Some(end) => body[..end].trim(),
None => body.trim(),
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn a_fenced_reply_is_unwrapped() {
assert_eq!(strip_code_fence("```json\n{\"a\": 1}\n```"), "{\"a\": 1}");
assert_eq!(strip_code_fence("```\n{\"a\": 1}\n```"), "{\"a\": 1}");
assert_eq!(strip_code_fence(" {\"a\": 1} "), "{\"a\": 1}");
assert_eq!(strip_code_fence("```json\n{\"a\": 1}"), "{\"a\": 1}");
}
#[test]
fn a_repair_names_the_reason_and_asks_for_only_json() {
let instruction = repair_instruction("the answer is not valid JSON");
assert!(instruction.contains("the answer is not valid JSON"));
assert!(instruction.contains("only a JSON value"));
}
}