Skip to main content

sim_lib_agent_runner_core/
output.rs

1use sim_kernel::{Expr, Symbol};
2use sim_shape::Shape;
3
4use crate::grammar::shape_to_grammar;
5
6/// Model-request extension key for the required return codec.
7pub const RETURN_CODEC_EXTRA: &str = "return-codec";
8/// Model-request extension key for the normalized return shape expression.
9pub const RETURN_SHAPE_EXTRA: &str = "return-shape";
10/// Model-request extension key for a best-effort output grammar.
11pub const OUTPUT_GRAMMAR_EXTRA: &str = "output-grammar";
12/// Model-request extension key marking whether the grammar is required.
13pub const OUTPUT_GRAMMAR_REQUIRED_EXTRA: &str = "output-grammar-required";
14
15/// Output obligations attached to a model request.
16#[derive(Clone, Debug, PartialEq, Eq)]
17pub struct OutputContract {
18    /// Codec the receiver must use for the return value.
19    pub codec: Symbol,
20    /// Normalized shape expression describing the return value.
21    pub shape_expr: Expr,
22    /// Optional constrained-generation grammar derived from the shape.
23    pub grammar: Option<String>,
24    /// Whether a runner that accepts grammars must enforce the grammar.
25    pub required: bool,
26}
27
28impl OutputContract {
29    /// Builds an output contract with an explicit grammar.
30    pub fn new(codec: Symbol, shape_expr: Expr, grammar: Option<String>, required: bool) -> Self {
31        Self {
32            codec,
33            shape_expr,
34            grammar,
35            required,
36        }
37    }
38
39    /// Builds an output contract, deriving a grammar when the shape supports it.
40    pub fn for_shape(codec: Symbol, shape_expr: Expr, shape: &dyn Shape, required: bool) -> Self {
41        Self::new(codec, shape_expr, shape_to_grammar(shape).ok(), required)
42    }
43
44    /// Appends this contract to a [`ModelRequest`](crate::ModelRequest) `extra` map.
45    pub fn into_extra_entries(self, extra: &mut Vec<(Expr, Expr)>) {
46        extra.push(entry(RETURN_CODEC_EXTRA, Expr::Symbol(self.codec)));
47        extra.push(entry(RETURN_SHAPE_EXTRA, self.shape_expr));
48        if let Some(grammar) = self.grammar {
49            extra.push(entry(OUTPUT_GRAMMAR_EXTRA, Expr::String(grammar)));
50        }
51        extra.push(entry(
52            OUTPUT_GRAMMAR_REQUIRED_EXTRA,
53            Expr::Bool(self.required),
54        ));
55    }
56}
57
58fn entry(name: &str, value: Expr) -> (Expr, Expr) {
59    (Expr::Symbol(Symbol::new(name)), value)
60}
61
62#[cfg(test)]
63mod tests {
64    use super::{
65        OUTPUT_GRAMMAR_EXTRA, OUTPUT_GRAMMAR_REQUIRED_EXTRA, OutputContract, RETURN_CODEC_EXTRA,
66        RETURN_SHAPE_EXTRA,
67    };
68    use sim_kernel::{
69        Cx, Expr, MatchScore, Result, ShapeDoc, ShapeMatch, Symbol, Value, shape::Shape,
70    };
71    use sim_shape::{ExprKind, ExprKindShape};
72
73    struct NamedRecordShape;
74
75    impl Shape for NamedRecordShape {
76        fn symbol(&self) -> Option<Symbol> {
77            Some(Symbol::qualified("bridge", "NamedRecord"))
78        }
79
80        fn check_value(&self, _cx: &mut Cx, _value: Value) -> Result<ShapeMatch> {
81            Ok(ShapeMatch::accept(MatchScore::exact(1)))
82        }
83
84        fn check_expr(&self, _cx: &mut Cx, _expr: &Expr) -> Result<ShapeMatch> {
85            Ok(ShapeMatch::accept(MatchScore::exact(1)))
86        }
87
88        fn describe(&self, _cx: &mut Cx) -> Result<ShapeDoc> {
89            Ok(ShapeDoc::new("named record"))
90        }
91    }
92
93    #[test]
94    fn supported_shape_attaches_grammar() {
95        let contract = OutputContract::for_shape(
96            Symbol::qualified("codec", "bridge"),
97            Expr::Symbol(Symbol::qualified("shape", "String")),
98            &ExprKindShape::new(ExprKind::String),
99            true,
100        );
101        let mut extra = Vec::new();
102        contract.into_extra_entries(&mut extra);
103
104        assert_eq!(
105            field(&extra, RETURN_CODEC_EXTRA),
106            Some(&Expr::Symbol(Symbol::qualified("codec", "bridge")))
107        );
108        assert_eq!(
109            field(&extra, RETURN_SHAPE_EXTRA),
110            Some(&Expr::Symbol(Symbol::qualified("shape", "String")))
111        );
112        assert_eq!(
113            field(&extra, OUTPUT_GRAMMAR_EXTRA),
114            Some(&Expr::String(r#"{"type":"string"}"#.to_owned()))
115        );
116        assert_eq!(
117            field(&extra, OUTPUT_GRAMMAR_REQUIRED_EXTRA),
118            Some(&Expr::Bool(true))
119        );
120    }
121
122    #[test]
123    fn named_record_shape_omits_grammar_without_error() {
124        let contract = OutputContract::for_shape(
125            Symbol::qualified("codec", "bridge"),
126            Expr::Symbol(Symbol::qualified("bridge", "NamedRecord")),
127            &NamedRecordShape,
128            true,
129        );
130        let mut extra = Vec::new();
131        contract.into_extra_entries(&mut extra);
132
133        assert!(field(&extra, OUTPUT_GRAMMAR_EXTRA).is_none());
134        assert_eq!(
135            field(&extra, RETURN_SHAPE_EXTRA),
136            Some(&Expr::Symbol(Symbol::qualified("bridge", "NamedRecord")))
137        );
138        assert_eq!(
139            field(&extra, OUTPUT_GRAMMAR_REQUIRED_EXTRA),
140            Some(&Expr::Bool(true))
141        );
142    }
143
144    fn field<'a>(entries: &'a [(Expr, Expr)], name: &str) -> Option<&'a Expr> {
145        entries.iter().find_map(|(key, value)| {
146            if *key == Expr::Symbol(Symbol::new(name)) {
147                Some(value)
148            } else {
149                None
150            }
151        })
152    }
153}