Skip to main content

nika_core/ast/
output.rs

1//! Output Policy - format and validation configuration
2//!
3//! Defines how task output should be formatted and validated:
4//! - `OutputFormat`: Text (default) or JSON
5//! - `SchemaRef`: Inline JSON Schema object or file path
6//! - `OutputPolicy`: Format + optional schema validation + retry config
7
8use serde::de::{self, Deserializer, MapAccess, Visitor};
9use serde::{Deserialize, Serialize};
10use serde_json::Value as JsonValue;
11use std::fmt;
12
13/// Reference to a JSON Schema - either inline or file path
14#[derive(Debug, Clone, Serialize)]
15#[serde(untagged)]
16pub enum SchemaRef {
17    /// Inline JSON Schema object
18    Inline(JsonValue),
19    /// Path to JSON Schema file
20    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/// Output policy configuration
65#[derive(Debug, Clone, Deserialize, Default)]
66pub struct OutputPolicy {
67    /// Output format (text or json)
68    #[serde(default)]
69    pub format: OutputFormat,
70
71    /// JSON Schema for output validation (inline object or file path)
72    #[serde(default)]
73    pub schema: Option<SchemaRef>,
74
75    /// JSON example — Nika auto-derives the JSON Schema from this at runtime.
76    /// Alternative to `schema`. When set, takes precedence over `schema`.
77    #[serde(default)]
78    pub from_example: Option<SchemaRef>,
79
80    /// Maximum retry attempts on validation failure (default: 2)
81    #[serde(default)]
82    pub max_retries: Option<u8>,
83
84    /// Original StructuredOutputSpec when this policy was bridged from `structured:`.
85    /// Preserves user's layer toggle config (enable_tool_injection, enable_repair, etc.)
86    /// that would otherwise be lost in the OutputPolicy→StructuredOutputSpec roundtrip.
87    #[serde(skip)]
88    pub source_structured_spec: Option<super::structured::StructuredOutputSpec>,
89}
90
91impl OutputPolicy {
92    /// Check if this policy requires structured output validation.
93    ///
94    /// Returns true when format is JSON and a schema is provided.
95    /// This is used by the executor to decide whether to use StructuredOutputEngine.
96    pub fn is_structured(&self) -> bool {
97        self.format == OutputFormat::Json && (self.schema.is_some() || self.from_example.is_some())
98    }
99
100    /// Convert to StructuredOutputSpec for use with StructuredOutputEngine.
101    ///
102    /// Returns None if this policy doesn't require structured output.
103    /// If this policy was bridged from a `structured:` config (via `to_output_policy()`),
104    /// returns the original spec with all user-configured layer toggles preserved.
105    /// Otherwise, constructs a spec with permissive defaults (all layers enabled).
106    pub fn to_structured_spec(&self) -> Option<super::structured::StructuredOutputSpec> {
107        if !self.is_structured() {
108            return None;
109        }
110
111        // Return original spec if available (preserves user's layer toggles)
112        if let Some(ref spec) = self.source_structured_spec {
113            return Some(spec.clone());
114        }
115
116        // Fallback: construct spec from OutputPolicy fields (output: block path)
117        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/// Output format enum
132#[derive(Debug, Clone, Deserialize, Default, PartialEq)]
133#[serde(rename_all = "lowercase")]
134pub enum OutputFormat {
135    /// Raw text output (default)
136    #[default]
137    Text,
138
139    /// JSON parsed output
140    Json,
141
142    /// YAML formatted output
143    Yaml,
144
145    /// Markdown formatted output
146    Markdown,
147
148    /// Binary (raw bytes, not text content)
149    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        // Verify schema content
195        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    // ========== is_structured() tests ==========
220
221    #[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        // Edge case: text format with schema should NOT be structured
260        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    // ========== to_structured_spec() tests ==========
270
271    #[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}