oxify_model/
json_schema.rs

1//! JSON Schema generation for workflow models
2//!
3//! This module provides functionality to generate JSON Schema documents
4//! from workflow models for validation, documentation, and integration.
5
6use crate::{NodeKind, Workflow};
7use serde::{Deserialize, Serialize};
8use serde_json::{json, Value};
9use std::collections::HashMap;
10
11#[cfg(feature = "openapi")]
12use utoipa::ToSchema;
13
14/// JSON Schema document
15#[derive(Debug, Clone, Serialize, Deserialize)]
16#[cfg_attr(feature = "openapi", derive(ToSchema))]
17pub struct JsonSchema {
18    /// Schema version (always "<https://json-schema.org/draft/2020-12/schema>")
19    #[serde(rename = "$schema")]
20    pub schema: String,
21
22    /// Schema identifier
23    #[serde(rename = "$id", skip_serializing_if = "Option::is_none")]
24    pub id: Option<String>,
25
26    /// Schema title
27    #[serde(skip_serializing_if = "Option::is_none")]
28    pub title: Option<String>,
29
30    /// Schema description
31    #[serde(skip_serializing_if = "Option::is_none")]
32    pub description: Option<String>,
33
34    /// Type of the schema
35    #[serde(rename = "type")]
36    pub schema_type: String,
37
38    /// Properties for object types
39    #[serde(skip_serializing_if = "Option::is_none")]
40    pub properties: Option<HashMap<String, JsonSchema>>,
41
42    /// Required property names
43    #[serde(skip_serializing_if = "Option::is_none")]
44    pub required: Option<Vec<String>>,
45
46    /// Items schema for array types
47    #[serde(skip_serializing_if = "Option::is_none")]
48    pub items: Option<Box<JsonSchema>>,
49
50    /// Enum values for enums
51    #[serde(rename = "enum", skip_serializing_if = "Option::is_none")]
52    pub enum_values: Option<Vec<Value>>,
53
54    /// Additional properties
55    #[serde(skip_serializing_if = "Option::is_none")]
56    pub additional_properties: Option<Box<JsonSchema>>,
57
58    /// Schema definitions
59    #[serde(rename = "$defs", skip_serializing_if = "Option::is_none")]
60    pub definitions: Option<HashMap<String, JsonSchema>>,
61
62    /// Minimum value for numbers
63    #[serde(skip_serializing_if = "Option::is_none")]
64    pub minimum: Option<f64>,
65
66    /// Maximum value for numbers
67    #[serde(skip_serializing_if = "Option::is_none")]
68    pub maximum: Option<f64>,
69
70    /// Pattern for strings
71    #[serde(skip_serializing_if = "Option::is_none")]
72    pub pattern: Option<String>,
73
74    /// Format for strings
75    #[serde(skip_serializing_if = "Option::is_none")]
76    pub format: Option<String>,
77
78    /// One of (union types)
79    #[serde(rename = "oneOf", skip_serializing_if = "Option::is_none")]
80    pub one_of: Option<Vec<JsonSchema>>,
81
82    /// Any of (union types with overlap)
83    #[serde(rename = "anyOf", skip_serializing_if = "Option::is_none")]
84    pub any_of: Option<Vec<JsonSchema>>,
85
86    /// All of (intersection types)
87    #[serde(rename = "allOf", skip_serializing_if = "Option::is_none")]
88    pub all_of: Option<Vec<JsonSchema>>,
89
90    /// Reference to another schema
91    #[serde(rename = "$ref", skip_serializing_if = "Option::is_none")]
92    pub reference: Option<String>,
93}
94
95impl JsonSchema {
96    /// Create a new empty schema
97    pub fn new(schema_type: &str) -> Self {
98        Self {
99            schema: "https://json-schema.org/draft/2020-12/schema".to_string(),
100            id: None,
101            title: None,
102            description: None,
103            schema_type: schema_type.to_string(),
104            properties: None,
105            required: None,
106            items: None,
107            enum_values: None,
108            additional_properties: None,
109            definitions: None,
110            minimum: None,
111            maximum: None,
112            pattern: None,
113            format: None,
114            one_of: None,
115            any_of: None,
116            all_of: None,
117            reference: None,
118        }
119    }
120
121    /// Create a string schema
122    pub fn string() -> Self {
123        Self::new("string")
124    }
125
126    /// Create a number schema
127    pub fn number() -> Self {
128        Self::new("number")
129    }
130
131    /// Create an integer schema
132    pub fn integer() -> Self {
133        Self::new("integer")
134    }
135
136    /// Create a boolean schema
137    pub fn boolean() -> Self {
138        Self::new("boolean")
139    }
140
141    /// Create an array schema
142    pub fn array(items: JsonSchema) -> Self {
143        let mut schema = Self::new("array");
144        schema.items = Some(Box::new(items));
145        schema
146    }
147
148    /// Create an object schema
149    pub fn object() -> Self {
150        let mut schema = Self::new("object");
151        schema.properties = Some(HashMap::new());
152        schema
153    }
154
155    /// Create a reference schema
156    pub fn reference(ref_path: &str) -> Self {
157        let mut schema = Self::new("object");
158        schema.reference = Some(ref_path.to_string());
159        schema
160    }
161
162    /// Set title
163    pub fn with_title(mut self, title: String) -> Self {
164        self.title = Some(title);
165        self
166    }
167
168    /// Set description
169    pub fn with_description(mut self, description: String) -> Self {
170        self.description = Some(description);
171        self
172    }
173
174    /// Set ID
175    pub fn with_id(mut self, id: String) -> Self {
176        self.id = Some(id);
177        self
178    }
179
180    /// Add a property to an object schema
181    pub fn add_property(&mut self, name: String, schema: JsonSchema) {
182        if self.properties.is_none() {
183            self.properties = Some(HashMap::new());
184        }
185        if let Some(props) = &mut self.properties {
186            props.insert(name, schema);
187        }
188    }
189
190    /// Add a required property
191    pub fn add_required(&mut self, name: String) {
192        if self.required.is_none() {
193            self.required = Some(Vec::new());
194        }
195        if let Some(req) = &mut self.required {
196            req.push(name);
197        }
198    }
199
200    /// Set enum values
201    pub fn with_enum(mut self, values: Vec<Value>) -> Self {
202        self.enum_values = Some(values);
203        self
204    }
205
206    /// Set pattern for strings
207    pub fn with_pattern(mut self, pattern: String) -> Self {
208        self.pattern = Some(pattern);
209        self
210    }
211
212    /// Set format for strings
213    pub fn with_format(mut self, format: String) -> Self {
214        self.format = Some(format);
215        self
216    }
217
218    /// Add a definition
219    pub fn add_definition(&mut self, name: String, schema: JsonSchema) {
220        if self.definitions.is_none() {
221            self.definitions = Some(HashMap::new());
222        }
223        if let Some(defs) = &mut self.definitions {
224            defs.insert(name, schema);
225        }
226    }
227}
228
229/// Schema generator for workflows
230pub struct WorkflowSchemaGenerator {
231    /// Include optional fields
232    pub include_optional: bool,
233
234    /// Include examples
235    pub include_examples: bool,
236
237    /// Include descriptions
238    pub include_descriptions: bool,
239}
240
241impl WorkflowSchemaGenerator {
242    /// Create a new schema generator with default settings
243    pub fn new() -> Self {
244        Self {
245            include_optional: true,
246            include_examples: false,
247            include_descriptions: true,
248        }
249    }
250
251    /// Generate JSON Schema for a workflow
252    pub fn generate_workflow_schema(&self) -> JsonSchema {
253        let mut schema = JsonSchema::object()
254            .with_id("https://oxify.dev/schemas/workflow.json".to_string())
255            .with_title("Workflow".to_string());
256
257        if self.include_descriptions {
258            schema =
259                schema.with_description("A workflow defining a sequence of operations".to_string());
260        }
261
262        // Add basic properties
263        schema.add_property(
264            "id".to_string(),
265            JsonSchema::string().with_format("uuid".to_string()),
266        );
267        schema.add_required("id".to_string());
268
269        schema.add_property("metadata".to_string(), self.generate_metadata_schema());
270        schema.add_required("metadata".to_string());
271
272        schema.add_property(
273            "nodes".to_string(),
274            JsonSchema::array(self.generate_node_schema()),
275        );
276        schema.add_required("nodes".to_string());
277
278        schema.add_property(
279            "edges".to_string(),
280            JsonSchema::array(self.generate_edge_schema()),
281        );
282        schema.add_required("edges".to_string());
283
284        // Add node type definitions
285        self.add_node_type_definitions(&mut schema);
286
287        schema
288    }
289
290    /// Generate schema for workflow metadata
291    fn generate_metadata_schema(&self) -> JsonSchema {
292        let mut schema = JsonSchema::object();
293
294        if self.include_descriptions {
295            schema = schema.with_description("Workflow metadata".to_string());
296        }
297
298        schema.add_property("name".to_string(), JsonSchema::string());
299        schema.add_required("name".to_string());
300
301        schema.add_property("description".to_string(), JsonSchema::string());
302
303        schema.add_property("version".to_string(), JsonSchema::string());
304        schema.add_required("version".to_string());
305
306        schema.add_property(
307            "created_at".to_string(),
308            JsonSchema::string().with_format("date-time".to_string()),
309        );
310        schema.add_required("created_at".to_string());
311
312        schema.add_property(
313            "updated_at".to_string(),
314            JsonSchema::string().with_format("date-time".to_string()),
315        );
316        schema.add_required("updated_at".to_string());
317
318        schema.add_property("tags".to_string(), JsonSchema::array(JsonSchema::string()));
319
320        schema
321    }
322
323    /// Generate schema for a node
324    fn generate_node_schema(&self) -> JsonSchema {
325        let mut schema = JsonSchema::object();
326
327        if self.include_descriptions {
328            schema = schema.with_description("A workflow node".to_string());
329        }
330
331        schema.add_property(
332            "id".to_string(),
333            JsonSchema::string().with_format("uuid".to_string()),
334        );
335        schema.add_required("id".to_string());
336
337        schema.add_property("name".to_string(), JsonSchema::string());
338        schema.add_required("name".to_string());
339
340        schema.add_property("kind".to_string(), self.generate_node_kind_schema());
341        schema.add_required("kind".to_string());
342
343        schema
344    }
345
346    /// Generate schema for node kind
347    fn generate_node_kind_schema(&self) -> JsonSchema {
348        JsonSchema::string().with_enum(vec![
349            json!("Start"),
350            json!("End"),
351            json!("LLM"),
352            json!("Retriever"),
353            json!("Code"),
354            json!("IfElse"),
355            json!("Tool"),
356            json!("Loop"),
357            json!("TryCatch"),
358            json!("SubWorkflow"),
359            json!("Switch"),
360            json!("Parallel"),
361            json!("Approval"),
362            json!("Form"),
363        ])
364    }
365
366    /// Generate schema for an edge
367    fn generate_edge_schema(&self) -> JsonSchema {
368        let mut schema = JsonSchema::object();
369
370        if self.include_descriptions {
371            schema = schema.with_description("A workflow edge connecting two nodes".to_string());
372        }
373
374        schema.add_property(
375            "id".to_string(),
376            JsonSchema::string().with_format("uuid".to_string()),
377        );
378        schema.add_required("id".to_string());
379
380        schema.add_property(
381            "from".to_string(),
382            JsonSchema::string().with_format("uuid".to_string()),
383        );
384        schema.add_required("from".to_string());
385
386        schema.add_property(
387            "to".to_string(),
388            JsonSchema::string().with_format("uuid".to_string()),
389        );
390        schema.add_required("to".to_string());
391
392        schema
393    }
394
395    /// Add node type definitions to schema
396    fn add_node_type_definitions(&self, schema: &mut JsonSchema) {
397        // LLM config
398        let mut llm_config = JsonSchema::object();
399        llm_config.add_property("model".to_string(), JsonSchema::string());
400        llm_config.add_property("prompt".to_string(), JsonSchema::string());
401        llm_config.add_property("temperature".to_string(), JsonSchema::number());
402        llm_config.add_property("max_tokens".to_string(), JsonSchema::integer());
403        schema.add_definition("LlmConfig".to_string(), llm_config);
404
405        // Script config
406        let mut script_config = JsonSchema::object();
407        script_config.add_property(
408            "language".to_string(),
409            JsonSchema::string().with_enum(vec![
410                json!("Python"),
411                json!("JavaScript"),
412                json!("TypeScript"),
413                json!("Bash"),
414            ]),
415        );
416        script_config.add_property("code".to_string(), JsonSchema::string());
417        schema.add_definition("ScriptConfig".to_string(), script_config);
418
419        // Condition
420        let mut condition = JsonSchema::object();
421        condition.add_property("expression".to_string(), JsonSchema::string());
422        schema.add_definition("Condition".to_string(), condition);
423    }
424
425    /// Generate schema for a specific node type
426    pub fn generate_node_type_schema(&self, node_kind: &NodeKind) -> JsonSchema {
427        match node_kind {
428            NodeKind::Start | NodeKind::End => {
429                JsonSchema::object().with_description(format!("{:?} node", node_kind))
430            }
431            NodeKind::LLM(_) => JsonSchema::reference("#/$defs/LlmConfig"),
432            NodeKind::Code(_) => JsonSchema::reference("#/$defs/ScriptConfig"),
433            _ => {
434                JsonSchema::object().with_description(format!("{:?} node configuration", node_kind))
435            }
436        }
437    }
438
439    /// Validate a workflow against the schema
440    pub fn validate_workflow(&self, workflow: &Workflow) -> Result<(), Vec<String>> {
441        let mut errors = Vec::new();
442
443        // Basic validation
444        if workflow.nodes.is_empty() {
445            errors.push("Workflow must have at least one node".to_string());
446        }
447
448        // Check for start node
449        if !workflow
450            .nodes
451            .iter()
452            .any(|n| matches!(n.kind, NodeKind::Start))
453        {
454            errors.push("Workflow must have a Start node".to_string());
455        }
456
457        // Check for end node
458        if !workflow
459            .nodes
460            .iter()
461            .any(|n| matches!(n.kind, NodeKind::End))
462        {
463            errors.push("Workflow must have an End node".to_string());
464        }
465
466        if errors.is_empty() {
467            Ok(())
468        } else {
469            Err(errors)
470        }
471    }
472}
473
474impl Default for WorkflowSchemaGenerator {
475    fn default() -> Self {
476        Self::new()
477    }
478}
479
480/// Generate a JSON Schema for a workflow
481pub fn generate_workflow_schema() -> JsonSchema {
482    WorkflowSchemaGenerator::new().generate_workflow_schema()
483}
484
485/// Export schema as JSON string
486pub fn schema_to_json(schema: &JsonSchema) -> Result<String, serde_json::Error> {
487    serde_json::to_string_pretty(schema)
488}
489
490/// Export schema as JSON value
491pub fn schema_to_value(schema: &JsonSchema) -> Result<Value, serde_json::Error> {
492    serde_json::to_value(schema)
493}
494
495#[cfg(test)]
496mod tests {
497    use super::*;
498    use crate::{Edge, Node, NodeKind};
499
500    #[test]
501    fn test_create_basic_schema() {
502        let schema = JsonSchema::string();
503        assert_eq!(schema.schema_type, "string");
504        assert_eq!(
505            schema.schema,
506            "https://json-schema.org/draft/2020-12/schema"
507        );
508    }
509
510    #[test]
511    fn test_create_object_schema() {
512        let mut schema = JsonSchema::object();
513        schema.add_property("name".to_string(), JsonSchema::string());
514        schema.add_required("name".to_string());
515
516        assert_eq!(schema.schema_type, "object");
517        assert!(schema.properties.is_some());
518        assert_eq!(schema.properties.as_ref().unwrap().len(), 1);
519        assert_eq!(schema.required.as_ref().unwrap().len(), 1);
520    }
521
522    #[test]
523    fn test_create_array_schema() {
524        let schema = JsonSchema::array(JsonSchema::string());
525        assert_eq!(schema.schema_type, "array");
526        assert!(schema.items.is_some());
527    }
528
529    #[test]
530    fn test_enum_schema() {
531        let schema = JsonSchema::string().with_enum(vec![
532            json!("option1"),
533            json!("option2"),
534            json!("option3"),
535        ]);
536
537        assert!(schema.enum_values.is_some());
538        assert_eq!(schema.enum_values.as_ref().unwrap().len(), 3);
539    }
540
541    #[test]
542    fn test_generate_workflow_schema() {
543        let generator = WorkflowSchemaGenerator::new();
544        let schema = generator.generate_workflow_schema();
545
546        assert_eq!(schema.schema_type, "object");
547        assert!(schema.properties.is_some());
548        assert!(schema.required.is_some());
549
550        let props = schema.properties.as_ref().unwrap();
551        assert!(props.contains_key("id"));
552        assert!(props.contains_key("metadata"));
553        assert!(props.contains_key("nodes"));
554        assert!(props.contains_key("edges"));
555
556        let required = schema.required.as_ref().unwrap();
557        assert!(required.contains(&"id".to_string()));
558        assert!(required.contains(&"metadata".to_string()));
559        assert!(required.contains(&"nodes".to_string()));
560        assert!(required.contains(&"edges".to_string()));
561    }
562
563    #[test]
564    fn test_schema_serialization() {
565        let schema = JsonSchema::string()
566            .with_title("Name".to_string())
567            .with_description("A person's name".to_string());
568
569        let json = schema_to_json(&schema).unwrap();
570        assert!(json.contains("Name"));
571        assert!(json.contains("person's name"));
572    }
573
574    #[test]
575    fn test_validate_workflow_missing_start() {
576        let mut workflow = Workflow::new("Test".to_string());
577        let end_node = Node::new("End".to_string(), NodeKind::End);
578        workflow.add_node(end_node);
579
580        let generator = WorkflowSchemaGenerator::new();
581        let result = generator.validate_workflow(&workflow);
582
583        assert!(result.is_err());
584        let errors = result.unwrap_err();
585        assert!(errors.iter().any(|e| e.contains("Start")));
586    }
587
588    #[test]
589    fn test_validate_workflow_missing_end() {
590        let mut workflow = Workflow::new("Test".to_string());
591        let start_node = Node::new("Start".to_string(), NodeKind::Start);
592        workflow.add_node(start_node);
593
594        let generator = WorkflowSchemaGenerator::new();
595        let result = generator.validate_workflow(&workflow);
596
597        assert!(result.is_err());
598        let errors = result.unwrap_err();
599        assert!(errors.iter().any(|e| e.contains("End")));
600    }
601
602    #[test]
603    fn test_validate_valid_workflow() {
604        let mut workflow = Workflow::new("Test".to_string());
605
606        let start_node = Node::new("Start".to_string(), NodeKind::Start);
607        let start_id = start_node.id;
608        workflow.add_node(start_node);
609
610        let end_node = Node::new("End".to_string(), NodeKind::End);
611        let end_id = end_node.id;
612        workflow.add_node(end_node);
613
614        workflow.add_edge(Edge::new(start_id, end_id));
615
616        let generator = WorkflowSchemaGenerator::new();
617        let result = generator.validate_workflow(&workflow);
618
619        assert!(result.is_ok());
620    }
621
622    #[test]
623    fn test_node_kind_schema() {
624        let generator = WorkflowSchemaGenerator::new();
625        let schema = generator.generate_node_kind_schema();
626
627        assert_eq!(schema.schema_type, "string");
628        assert!(schema.enum_values.is_some());
629
630        let enums = schema.enum_values.as_ref().unwrap();
631        assert!(enums.contains(&json!("Start")));
632        assert!(enums.contains(&json!("End")));
633        assert!(enums.contains(&json!("LLM")));
634    }
635
636    #[test]
637    fn test_metadata_schema() {
638        let generator = WorkflowSchemaGenerator::new();
639        let schema = generator.generate_metadata_schema();
640
641        assert_eq!(schema.schema_type, "object");
642        let props = schema.properties.as_ref().unwrap();
643        assert!(props.contains_key("name"));
644        assert!(props.contains_key("version"));
645        assert!(props.contains_key("created_at"));
646        assert!(props.contains_key("updated_at"));
647    }
648
649    #[test]
650    fn test_reference_schema() {
651        let schema = JsonSchema::reference("#/$defs/LlmConfig");
652        assert!(schema.reference.is_some());
653        assert_eq!(schema.reference.unwrap(), "#/$defs/LlmConfig");
654    }
655
656    #[test]
657    fn test_schema_with_pattern() {
658        let schema = JsonSchema::string().with_pattern("^[a-zA-Z0-9_-]+$".to_string());
659
660        assert!(schema.pattern.is_some());
661        assert_eq!(schema.pattern.unwrap(), "^[a-zA-Z0-9_-]+$");
662    }
663
664    #[test]
665    fn test_schema_with_format() {
666        let schema = JsonSchema::string().with_format("uuid".to_string());
667
668        assert!(schema.format.is_some());
669        assert_eq!(schema.format.unwrap(), "uuid");
670    }
671
672    #[test]
673    fn test_schema_definitions() {
674        let mut schema = JsonSchema::object();
675        schema.add_definition("CustomType".to_string(), JsonSchema::string());
676
677        assert!(schema.definitions.is_some());
678        assert_eq!(schema.definitions.as_ref().unwrap().len(), 1);
679    }
680}
oxify_model/json_schema.rs

oxify_model/
json_schema.rs
