sim_lib_agent_runner_core/
output.rs

1use sim_kernel::{Expr, Symbol};
2use sim_shape::{GrammarDialect, GrammarGraph, Shape, shape_grammar_graph, shape_json_schema};
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 for the concrete grammar dialect.
13pub const OUTPUT_GRAMMAR_DIALECT_EXTRA: &str = "output-grammar-dialect";
14/// Model-request extension key marking whether the grammar is required.
15pub const OUTPUT_GRAMMAR_REQUIRED_EXTRA: &str = "output-grammar-required";
16
17/// Output obligations attached to a model request.
18#[derive(Clone, Debug, PartialEq, Eq)]
19pub struct OutputContract {
20    /// Codec the receiver must use for the return value.
21    pub codec: Symbol,
22    /// Normalized shape expression describing the return value.
23    pub shape_expr: Expr,
24    /// Optional constrained-generation grammar derived from the shape.
25    pub grammar: Option<String>,
26    /// Dialect used by [`Self::grammar`].
27    pub grammar_dialect: Option<GrammarDialect>,
28    /// Codec-neutral grammar source for provider-specific rendering.
29    pub grammar_graph: Option<GrammarGraph>,
30    /// Whether a runner that accepts grammars must enforce the grammar.
31    pub required: bool,
32}
33
34impl OutputContract {
35    /// Builds an output contract with an explicit grammar.
36    pub fn new(codec: Symbol, shape_expr: Expr, grammar: Option<String>, required: bool) -> Self {
37        let grammar_dialect = grammar.as_ref().map(|_| GrammarDialect::JsonSchema);
38        Self {
39            codec,
40            shape_expr,
41            grammar,
42            grammar_dialect,
43            grammar_graph: None,
44            required,
45        }
46    }
47
48    /// Builds an output contract, deriving a grammar when the shape supports it.
49    pub fn for_shape(codec: Symbol, shape_expr: Expr, shape: &dyn Shape, required: bool) -> Self {
50        let grammar_graph = shape_grammar_graph(shape).ok();
51        let grammar = grammar_graph
52            .as_ref()
53            .and_then(|_| shape_json_schema(shape).ok())
54            .or_else(|| shape_to_grammar(shape).ok());
55        let grammar_dialect = grammar.as_ref().map(|_| GrammarDialect::JsonSchema);
56        Self {
57            codec,
58            shape_expr,
59            grammar,
60            grammar_dialect,
61            grammar_graph,
62            required,
63        }
64    }
65
66    /// Appends this contract to a [`ModelRequest`](crate::ModelRequest) `extra` map.
67    pub fn into_extra_entries(self, extra: &mut Vec<(Expr, Expr)>) {
68        extra.push(entry(RETURN_CODEC_EXTRA, Expr::Symbol(self.codec)));
69        extra.push(entry(RETURN_SHAPE_EXTRA, self.shape_expr));
70        if let Some(grammar) = self.grammar {
71            extra.push(entry(OUTPUT_GRAMMAR_EXTRA, Expr::String(grammar)));
72        }
73        if let Some(dialect) = self.grammar_dialect {
74            extra.push(entry(
75                OUTPUT_GRAMMAR_DIALECT_EXTRA,
76                Expr::Symbol(grammar_dialect_symbol(dialect)),
77            ));
78        }
79        extra.push(entry(
80            OUTPUT_GRAMMAR_REQUIRED_EXTRA,
81            Expr::Bool(self.required),
82        ));
83    }
84}
85
86/// Returns the extension symbol used to name a grammar dialect.
87pub fn grammar_dialect_symbol(dialect: GrammarDialect) -> Symbol {
88    match dialect {
89        GrammarDialect::JsonSchema => Symbol::new("json-schema"),
90        GrammarDialect::Gbnf => Symbol::new("gbnf"),
91        GrammarDialect::SExpr => Symbol::new("sexpr"),
92    }
93}
94
95/// Parses a grammar dialect extension symbol.
96pub fn grammar_dialect_from_symbol(symbol: &Symbol) -> Option<GrammarDialect> {
97    match symbol.name.as_ref() {
98        "json-schema" if symbol.namespace.is_none() => Some(GrammarDialect::JsonSchema),
99        "gbnf" if symbol.namespace.is_none() => Some(GrammarDialect::Gbnf),
100        "sexpr" if symbol.namespace.is_none() => Some(GrammarDialect::SExpr),
101        _ => None,
102    }
103}
104
105fn entry(name: &str, value: Expr) -> (Expr, Expr) {
106    (Expr::Symbol(Symbol::new(name)), value)
107}
108
109#[cfg(test)]
110mod tests {
111    use super::{
112        OUTPUT_GRAMMAR_DIALECT_EXTRA, OUTPUT_GRAMMAR_EXTRA, OUTPUT_GRAMMAR_REQUIRED_EXTRA,
113        OutputContract, RETURN_CODEC_EXTRA, RETURN_SHAPE_EXTRA,
114    };
115    use sim_kernel::{
116        Cx, Expr, MatchScore, Result, ShapeDoc, ShapeMatch, Symbol, Value, shape::Shape,
117    };
118    use sim_shape::{ExprKind, ExprKindShape};
119
120    struct NamedRecordShape;
121
122    impl Shape for NamedRecordShape {
123        fn symbol(&self) -> Option<Symbol> {
124            Some(Symbol::qualified("bridge", "NamedRecord"))
125        }
126
127        fn check_value(&self, _cx: &mut Cx, _value: Value) -> Result<ShapeMatch> {
128            Ok(ShapeMatch::accept(MatchScore::exact(1)))
129        }
130
131        fn check_expr(&self, _cx: &mut Cx, _expr: &Expr) -> Result<ShapeMatch> {
132            Ok(ShapeMatch::accept(MatchScore::exact(1)))
133        }
134
135        fn describe(&self, _cx: &mut Cx) -> Result<ShapeDoc> {
136            Ok(ShapeDoc::new("named record"))
137        }
138    }
139
140    #[test]
141    fn supported_shape_attaches_grammar() {
142        let contract = OutputContract::for_shape(
143            Symbol::qualified("codec", "bridge"),
144            Expr::Symbol(Symbol::qualified("shape", "String")),
145            &ExprKindShape::new(ExprKind::String),
146            true,
147        );
148        let mut extra = Vec::new();
149        contract.into_extra_entries(&mut extra);
150
151        assert_eq!(
152            field(&extra, RETURN_CODEC_EXTRA),
153            Some(&Expr::Symbol(Symbol::qualified("codec", "bridge")))
154        );
155        assert_eq!(
156            field(&extra, RETURN_SHAPE_EXTRA),
157            Some(&Expr::Symbol(Symbol::qualified("shape", "String")))
158        );
159        assert_eq!(
160            field(&extra, OUTPUT_GRAMMAR_EXTRA),
161            Some(&Expr::String(r#"{"type":"string"}"#.to_owned()))
162        );
163        assert_eq!(
164            field(&extra, OUTPUT_GRAMMAR_DIALECT_EXTRA),
165            Some(&Expr::Symbol(Symbol::new("json-schema")))
166        );
167        assert_eq!(
168            field(&extra, OUTPUT_GRAMMAR_REQUIRED_EXTRA),
169            Some(&Expr::Bool(true))
170        );
171    }
172
173    #[test]
174    fn named_record_shape_omits_grammar_without_error() {
175        let contract = OutputContract::for_shape(
176            Symbol::qualified("codec", "bridge"),
177            Expr::Symbol(Symbol::qualified("bridge", "NamedRecord")),
178            &NamedRecordShape,
179            true,
180        );
181        let mut extra = Vec::new();
182        contract.into_extra_entries(&mut extra);
183
184        assert!(field(&extra, OUTPUT_GRAMMAR_EXTRA).is_none());
185        assert!(field(&extra, OUTPUT_GRAMMAR_DIALECT_EXTRA).is_none());
186        assert_eq!(
187            field(&extra, RETURN_SHAPE_EXTRA),
188            Some(&Expr::Symbol(Symbol::qualified("bridge", "NamedRecord")))
189        );
190        assert_eq!(
191            field(&extra, OUTPUT_GRAMMAR_REQUIRED_EXTRA),
192            Some(&Expr::Bool(true))
193        );
194    }
195
196    fn field<'a>(entries: &'a [(Expr, Expr)], name: &str) -> Option<&'a Expr> {
197        entries.iter().find_map(|(key, value)| {
198            if *key == Expr::Symbol(Symbol::new(name)) {
199                Some(value)
200            } else {
201                None
202            }
203        })
204    }
205}
sim_lib_agent_runner_core/output.rs

sim_lib_agent_runner_core/
output.rs
