1use serde::de::{self, Deserializer, MapAccess, Visitor};
9use serde::{Deserialize, Serialize};
10use serde_json::Value as JsonValue;
11use std::fmt;
12
13#[derive(Debug, Clone, Serialize)]
15#[serde(untagged)]
16pub enum SchemaRef {
17 Inline(JsonValue),
19 File(String),
21}
22
23impl<'de> Deserialize<'de> for SchemaRef {
24 fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
25 where
26 D: Deserializer<'de>,
27 {
28 struct SchemaRefVisitor;
29
30 impl<'de> Visitor<'de> for SchemaRefVisitor {
31 type Value = SchemaRef;
32
33 fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
34 formatter.write_str("a JSON Schema object or a file path string")
35 }
36
37 fn visit_str<E>(self, v: &str) -> Result<Self::Value, E>
38 where
39 E: de::Error,
40 {
41 Ok(SchemaRef::File(v.to_string()))
42 }
43
44 fn visit_string<E>(self, v: String) -> Result<Self::Value, E>
45 where
46 E: de::Error,
47 {
48 Ok(SchemaRef::File(v))
49 }
50
51 fn visit_map<A>(self, map: A) -> Result<Self::Value, A::Error>
52 where
53 A: MapAccess<'de>,
54 {
55 let value = JsonValue::deserialize(de::value::MapAccessDeserializer::new(map))?;
56 Ok(SchemaRef::Inline(value))
57 }
58 }
59
60 deserializer.deserialize_any(SchemaRefVisitor)
61 }
62}
63
64#[derive(Debug, Clone, Deserialize, Default)]
66pub struct OutputPolicy {
67 #[serde(default)]
69 pub format: OutputFormat,
70
71 #[serde(default)]
73 pub schema: Option<SchemaRef>,
74
75 #[serde(default)]
78 pub from_example: Option<SchemaRef>,
79
80 #[serde(default)]
82 pub max_retries: Option<u8>,
83
84 #[serde(skip)]
88 pub source_structured_spec: Option<super::structured::StructuredOutputSpec>,
89}
90
91impl OutputPolicy {
92 pub fn is_structured(&self) -> bool {
97 self.format == OutputFormat::Json && (self.schema.is_some() || self.from_example.is_some())
98 }
99
100 pub fn to_structured_spec(&self) -> Option<super::structured::StructuredOutputSpec> {
107 if !self.is_structured() {
108 return None;
109 }
110
111 if let Some(ref spec) = self.source_structured_spec {
113 return Some(spec.clone());
114 }
115
116 Some(super::structured::StructuredOutputSpec {
118 schema: self.schema.clone(),
119 from_example: self.from_example.clone(),
120 enable_extractor: None,
121 enable_tool_injection: None,
122 enable_retry: Some(true),
123 enable_repair: Some(true),
124 max_retries: self.max_retries,
125 repair_model: None,
126 strict: None,
127 })
128 }
129}
130
131#[derive(Debug, Clone, Deserialize, Default, PartialEq)]
133#[serde(rename_all = "lowercase")]
134pub enum OutputFormat {
135 #[default]
137 Text,
138
139 Json,
141
142 Yaml,
144
145 Markdown,
147
148 Binary,
150}
151
152#[cfg(test)]
153mod tests {
154 use super::*;
155 use crate::serde_yaml;
156
157 #[test]
158 fn parse_text_format() {
159 let yaml = "format: text";
160 let policy: OutputPolicy = serde_yaml::from_str(yaml).unwrap();
161 assert_eq!(policy.format, OutputFormat::Text);
162 assert!(policy.schema.is_none());
163 }
164
165 #[test]
166 fn parse_json_with_schema_file() {
167 let yaml = r#"
168 format: json
169 schema: .nika/schemas/result.json
170 "#;
171 let policy: OutputPolicy = serde_yaml::from_str(yaml).unwrap();
172 assert_eq!(policy.format, OutputFormat::Json);
173 assert!(
174 matches!(policy.schema, Some(SchemaRef::File(ref p)) if p == ".nika/schemas/result.json")
175 );
176 }
177
178 #[test]
179 fn parse_json_with_inline_schema() {
180 let yaml = r#"
181format: json
182schema:
183 type: object
184 properties:
185 name:
186 type: string
187 required:
188 - name
189"#;
190 let policy: OutputPolicy = serde_yaml::from_str(yaml).unwrap();
191 assert_eq!(policy.format, OutputFormat::Json);
192 assert!(matches!(policy.schema, Some(SchemaRef::Inline(_))));
193
194 if let Some(SchemaRef::Inline(schema)) = &policy.schema {
196 assert_eq!(schema["type"], "object");
197 assert!(schema["properties"]["name"].is_object());
198 }
199 }
200
201 #[test]
202 fn parse_max_retries() {
203 let yaml = r#"
204format: json
205max_retries: 3
206"#;
207 let policy: OutputPolicy = serde_yaml::from_str(yaml).unwrap();
208 assert_eq!(policy.max_retries, Some(3));
209 }
210
211 #[test]
212 fn default_is_text() {
213 let policy = OutputPolicy::default();
214 assert_eq!(policy.format, OutputFormat::Text);
215 assert!(policy.schema.is_none());
216 assert!(policy.max_retries.is_none());
217 }
218
219 #[test]
222 fn is_structured_true_when_json_with_schema() {
223 let yaml = r#"
224format: json
225schema:
226 type: object
227"#;
228 let policy: OutputPolicy = serde_yaml::from_str(yaml).unwrap();
229 assert!(policy.is_structured());
230 }
231
232 #[test]
233 fn is_structured_false_when_text() {
234 let policy = OutputPolicy::default();
235 assert!(!policy.is_structured());
236 }
237
238 #[test]
239 fn is_structured_false_when_json_without_schema() {
240 let yaml = "format: json";
241 let policy: OutputPolicy = serde_yaml::from_str(yaml).unwrap();
242 assert!(!policy.is_structured());
243 }
244
245 #[test]
246 fn is_structured_true_when_json_with_from_example() {
247 let yaml = r#"
248format: json
249from_example:
250 name: "Alice"
251 age: 30
252"#;
253 let policy: OutputPolicy = serde_yaml::from_str(yaml).unwrap();
254 assert!(policy.is_structured());
255 }
256
257 #[test]
258 fn is_structured_false_when_text_with_schema() {
259 let yaml = r#"
261format: text
262schema:
263 type: object
264"#;
265 let policy: OutputPolicy = serde_yaml::from_str(yaml).unwrap();
266 assert!(!policy.is_structured());
267 }
268
269 #[test]
272 fn to_structured_spec_returns_spec_when_structured() {
273 let yaml = r#"
274format: json
275schema:
276 type: object
277 properties:
278 name:
279 type: string
280max_retries: 5
281"#;
282 let policy: OutputPolicy = serde_yaml::from_str(yaml).unwrap();
283 let spec = policy.to_structured_spec();
284 assert!(spec.is_some());
285
286 let spec = spec.unwrap();
287 assert!(matches!(spec.schema, Some(SchemaRef::Inline(_))));
288 assert_eq!(spec.max_retries, Some(5));
289 assert_eq!(spec.enable_retry, Some(true));
290 assert_eq!(spec.enable_repair, Some(true));
291 }
292
293 #[test]
294 fn to_structured_spec_returns_none_when_not_structured() {
295 let policy = OutputPolicy::default();
296 assert!(policy.to_structured_spec().is_none());
297 }
298
299 #[test]
300 fn to_structured_spec_with_file_schema() {
301 let yaml = r#"
302format: json
303schema: ./schemas/user.json
304"#;
305 let policy: OutputPolicy = serde_yaml::from_str(yaml).unwrap();
306 let spec = policy.to_structured_spec();
307 assert!(spec.is_some());
308
309 let spec = spec.unwrap();
310 assert!(matches!(spec.schema, Some(SchemaRef::File(ref p)) if p == "./schemas/user.json"));
311 }
312}