@@ -4,15 +4,20 @@ use bamltype::jsonish::BamlValueWithFlags;
44use bamltype:: jsonish:: deserializer:: coercer:: run_user_checks;
55use bamltype:: jsonish:: deserializer:: deserialize_flags:: DeserializerConditions ;
66use indexmap:: IndexMap ;
7+ use minijinja:: UndefinedBehavior ;
8+ use minijinja:: value:: { Kwargs , Value as MiniJinjaValue } ;
79use regex:: Regex ;
8- use std:: sync:: LazyLock ;
10+ use serde_json:: { Value , json} ;
11+ use std:: collections:: HashMap ;
12+ use std:: sync:: { LazyLock , Mutex } ;
913use tracing:: { debug, trace} ;
1014
1115use super :: Adapter ;
1216use crate :: CallMetadata ;
1317use crate :: {
14- BamlType , BamlValue , ConstraintLevel , ConstraintResult , FieldMeta , Flag , JsonishError , Message ,
15- OutputFormatContent , ParseError , PredictError , Predicted , RenderOptions , Signature , TypeIR ,
18+ BamlType , BamlValue , ConstraintLevel , ConstraintResult , FieldMeta , Flag , InputRenderSpec ,
19+ JsonishError , Message , OutputFormatContent , ParseError , PredictError , Predicted , RenderOptions ,
20+ Signature , TypeIR ,
1621} ;
1722
1823/// Builds prompts and parses responses using the `[[ ## field ## ]]` delimiter protocol.
@@ -33,6 +38,101 @@ pub struct ChatAdapter;
3338static FIELD_HEADER_PATTERN : LazyLock < Regex > =
3439 LazyLock :: new ( || Regex :: new ( r"^\[\[ ## ([^#]+?) ## \]\]" ) . unwrap ( ) ) ;
3540
41+ const INPUT_RENDER_TEMPLATE_NAME : & str = "__input_field__" ;
42+
43+ #[ derive( Clone ) ]
44+ struct CachedInputRenderTemplate {
45+ env : minijinja:: Environment < ' static > ,
46+ }
47+
48+ static INPUT_RENDER_TEMPLATE_CACHE : LazyLock <
49+ Mutex < HashMap < & ' static str , CachedInputRenderTemplate > > ,
50+ > = LazyLock :: new ( || Mutex :: new ( HashMap :: new ( ) ) ) ;
51+
52+ fn regex_match ( value : String , regex : String ) -> bool {
53+ match Regex :: new ( & regex) {
54+ Ok ( re) => re. is_match ( & value) ,
55+ Err ( _) => false ,
56+ }
57+ }
58+
59+ fn sum_filter ( value : Vec < MiniJinjaValue > ) -> MiniJinjaValue {
60+ let int_sum: Option < i64 > = value
61+ . iter ( )
62+ . map ( |value| <i64 >:: try_from ( value. clone ( ) ) . ok ( ) )
63+ . collect :: < Option < Vec < _ > > > ( )
64+ . map ( |ints| ints. into_iter ( ) . sum ( ) ) ;
65+ let float_sum: Option < f64 > = value
66+ . into_iter ( )
67+ . map ( |value| <f64 >:: try_from ( value) . ok ( ) )
68+ . collect :: < Option < Vec < _ > > > ( )
69+ . map ( |floats| floats. into_iter ( ) . sum ( ) ) ;
70+ int_sum. map_or (
71+ float_sum. map_or ( MiniJinjaValue :: from ( 0 ) , MiniJinjaValue :: from) ,
72+ MiniJinjaValue :: from,
73+ )
74+ }
75+
76+ fn truncate_filter (
77+ value : String ,
78+ positional_length : Option < usize > ,
79+ kwargs : Kwargs ,
80+ ) -> Result < String , minijinja:: Error > {
81+ let kwarg_length: Option < usize > = kwargs. get ( "length" ) ?;
82+ let length = kwarg_length. or ( positional_length) . unwrap_or ( 255 ) ;
83+ let killwords: Option < bool > = kwargs. get ( "killwords" ) ?;
84+ let leeway: Option < usize > = kwargs. get ( "leeway" ) ?;
85+ let end: Option < String > = kwargs. get ( "end" ) ?;
86+ kwargs. assert_all_used ( ) ?;
87+
88+ let killwords = killwords. unwrap_or ( false ) ;
89+ let leeway = leeway. unwrap_or ( 5 ) ;
90+ let end = end. unwrap_or_else ( || "..." . to_string ( ) ) ;
91+ let value_len = value. chars ( ) . count ( ) ;
92+
93+ if value_len <= length. saturating_add ( leeway) {
94+ return Ok ( value) ;
95+ }
96+
97+ let trim_to = length. saturating_sub ( end. chars ( ) . count ( ) ) ;
98+ if trim_to == 0 {
99+ return Ok ( end. chars ( ) . take ( length) . collect ( ) ) ;
100+ }
101+
102+ let mut truncated: String = value. chars ( ) . take ( trim_to) . collect ( ) ;
103+ if !killwords {
104+ if let Some ( index) = truncated. rfind ( char:: is_whitespace) {
105+ if index > 0 {
106+ truncated. truncate ( index) ;
107+ }
108+ }
109+ truncated = truncated. trim_end ( ) . to_string ( ) ;
110+ }
111+
112+ Ok ( format ! ( "{truncated}{end}" ) )
113+ }
114+
115+ fn build_input_render_environment ( ) -> minijinja:: Environment < ' static > {
116+ // Keep this setup aligned with BAML's jinja env defaults, then add contrib filters.
117+ let mut env = minijinja:: Environment :: new ( ) ;
118+ env. set_formatter ( |output, state, value| {
119+ let value = if value. is_none ( ) {
120+ & MiniJinjaValue :: from ( "null" )
121+ } else {
122+ value
123+ } ;
124+ minijinja:: escape_formatter ( output, state, value)
125+ } ) ;
126+ env. set_debug ( true ) ;
127+ env. set_trim_blocks ( true ) ;
128+ env. set_lstrip_blocks ( true ) ;
129+ env. set_undefined_behavior ( UndefinedBehavior :: Strict ) ;
130+ env. add_filter ( "regex_match" , regex_match) ;
131+ env. add_filter ( "sum" , sum_filter) ;
132+ env. add_filter ( "truncate" , truncate_filter) ;
133+ env
134+ }
135+
36136fn render_field_type_schema (
37137 parent_format : & OutputFormatContent ,
38138 type_ir : & TypeIR ,
@@ -396,15 +496,19 @@ impl ChatAdapter {
396496 {
397497 let baml_value = input. to_baml_value ( ) ;
398498 let input_output_format = <I as BamlType >:: baml_output_format ( ) ;
499+ let input_json = build_input_context_value ( schema, & baml_value) ;
500+ let vars = Value :: Object ( serde_json:: Map :: new ( ) ) ;
399501
400502 let mut result = String :: new ( ) ;
401503 for field_spec in schema. input_fields ( ) {
402504 if let Some ( value) = value_for_path_relaxed ( & baml_value, field_spec. path ( ) ) {
403505 result. push_str ( & format ! ( "[[ ## {} ## ]]\n " , field_spec. lm_name) ) ;
404- result. push_str ( & format_baml_value_for_prompt_typed (
506+ result. push_str ( & render_input_field (
507+ field_spec,
405508 value,
509+ & input_json,
406510 input_output_format,
407- field_spec . format ,
511+ & vars ,
408512 ) ) ;
409513 result. push_str ( "\n \n " ) ;
410514 }
@@ -865,23 +969,117 @@ fn format_baml_value_for_prompt(value: &BamlValue) -> String {
865969 }
866970}
867971
868- fn format_baml_value_for_prompt_typed (
972+ fn render_input_field (
973+ field_spec : & crate :: FieldSchema ,
869974 value : & BamlValue ,
975+ input : & Value ,
870976 output_format : & OutputFormatContent ,
871- format : Option < & str > ,
977+ vars : & Value ,
872978) -> String {
873- let format = match format {
874- Some ( format) => format,
875- None => {
876- if let BamlValue :: String ( s) = value {
877- return s. clone ( ) ;
878- }
879- "json"
979+ match field_spec. input_render {
980+ InputRenderSpec :: Default => match value {
981+ BamlValue :: String ( s) => s. clone ( ) ,
982+ _ => bamltype:: internal_baml_jinja:: format_baml_value ( value, output_format, "json" )
983+ . unwrap_or_else ( |_| "<error>" . to_string ( ) ) ,
984+ } ,
985+ InputRenderSpec :: Format ( format) => {
986+ bamltype:: internal_baml_jinja:: format_baml_value ( value, output_format, format)
987+ . unwrap_or_else ( |_| "<error>" . to_string ( ) )
988+ }
989+ InputRenderSpec :: Jinja ( template) => {
990+ render_input_field_jinja ( template, field_spec, value, input, output_format, vars)
991+ }
992+ }
993+ }
994+
995+ fn build_input_context_value ( schema : & crate :: SignatureSchema , root : & BamlValue ) -> Value {
996+ let mut input_json = baml_value_to_render_json ( root) ;
997+ let Some ( root_map) = input_json. as_object_mut ( ) else {
998+ return input_json;
999+ } ;
1000+
1001+ // Provide alias lookups for top-level fields so templates can use either
1002+ // Rust field names (`input.question`) or prompt aliases (`input.query`).
1003+ for field_spec in schema. input_fields ( ) {
1004+ if field_spec. rust_name . contains ( '.' ) || field_spec. lm_name == field_spec. rust_name {
1005+ continue ;
1006+ }
1007+ if field_spec. path ( ) . iter ( ) . nth ( 1 ) . is_some ( ) {
1008+ continue ;
1009+ }
1010+ if let Some ( value) = root_map. get ( field_spec. rust_name . as_str ( ) ) . cloned ( ) {
1011+ root_map
1012+ . entry ( field_spec. lm_name . to_string ( ) )
1013+ . or_insert ( value) ;
8801014 }
1015+ }
1016+
1017+ input_json
1018+ }
1019+
1020+ fn baml_value_to_render_json ( value : & BamlValue ) -> Value {
1021+ serde_json:: to_value ( value) . unwrap_or ( Value :: Null )
1022+ }
1023+
1024+ fn render_input_field_jinja (
1025+ template : & ' static str ,
1026+ field_spec : & crate :: FieldSchema ,
1027+ value : & BamlValue ,
1028+ input : & Value ,
1029+ _output_format : & OutputFormatContent ,
1030+ vars : & Value ,
1031+ ) -> String {
1032+ let env = {
1033+ let mut cache = INPUT_RENDER_TEMPLATE_CACHE
1034+ . lock ( )
1035+ . expect ( "input render template cache lock poisoned" ) ;
1036+ cache
1037+ . entry ( template)
1038+ . or_insert_with ( || {
1039+ let mut env = build_input_render_environment ( ) ;
1040+ env. add_template ( INPUT_RENDER_TEMPLATE_NAME , template)
1041+ . unwrap_or_else ( |err| {
1042+ panic ! (
1043+ "failed to compile cached input render template for `{}` ({}): {err}" ,
1044+ field_spec. lm_name, field_spec. rust_name
1045+ )
1046+ } ) ;
1047+ CachedInputRenderTemplate { env }
1048+ } )
1049+ . env
1050+ . clone ( )
8811051 } ;
8821052
883- bamltype:: internal_baml_jinja:: format_baml_value ( value, output_format, format)
884- . unwrap_or_else ( |_| "<error>" . to_string ( ) )
1053+ let compiled = env
1054+ . get_template ( INPUT_RENDER_TEMPLATE_NAME )
1055+ . unwrap_or_else ( |err| {
1056+ panic ! (
1057+ "failed to fetch cached input render template for `{}` ({}): {err}" ,
1058+ field_spec. lm_name, field_spec. rust_name
1059+ )
1060+ } ) ;
1061+
1062+ let this = baml_value_to_render_json ( value) ;
1063+ let field = json ! ( {
1064+ "name" : field_spec. lm_name,
1065+ "rust_name" : field_spec. rust_name,
1066+ "type" : field_spec. type_ir. diagnostic_repr( ) . to_string( ) ,
1067+ } ) ;
1068+ let context = json ! ( {
1069+ "this" : this,
1070+ "input" : input,
1071+ "field" : field,
1072+ "vars" : vars,
1073+ } ) ;
1074+
1075+ compiled
1076+ . render ( minijinja:: Value :: from_serialize ( context) )
1077+ . unwrap_or_else ( |err| {
1078+ panic ! (
1079+ "failed to render input field `{}` (rust `{}`) with #[render(jinja = ...)] template `{}`: {err}" ,
1080+ field_spec. lm_name, field_spec. rust_name, template
1081+ )
1082+ } )
8851083}
8861084
8871085fn collect_flags_recursive ( value : & BamlValueWithFlags , flags : & mut Vec < Flag > ) {
0 commit comments