add input field render facet with jinja templates

Author: darinkishore

Cargo.lock

@@ -43,6 +43,7 @@ rig-core = { git = "https://github.com/0xPlaygrounds/rig", rev="e7849df" }
 enum_dispatch = "0.3.13"
 tracing = "0.1.44"
 tracing-subscriber = { version = "0.3.22", features = ["env-filter", "fmt"] }
+minijinja = { git = "https://github.com/boundaryml/minijinja.git", branch = "main", default-features = false, features = ["builtins", "serde"] }
 
 [package.metadata.cargo-machete]
 ignored = ["rig-core"]

crates/dspy-rs/Cargo.toml

@@ -1,5 +1,6 @@
 use anyhow::Result;
 use indexmap::IndexMap;
+use minijinja::UndefinedBehavior;
 use regex::Regex;
 use rig::tool::ToolDyn;
 use serde_json::{Value, json};
@@ -17,8 +18,8 @@ use crate::serde_utils::get_iter_from_value;
 use crate::utils::cache::CacheEntry;
 use crate::{
     BamlValue, Cache, Chat, ConstraintLevel, ConstraintResult, Example, FieldMeta, Flag,
-    JsonishError, LM, Message, MetaSignature, OutputFormatContent, ParseError, Prediction,
-    RenderOptions, Signature, TypeIR,
+    InputRenderSpec, JsonishError, LM, Message, MetaSignature, OutputFormatContent, ParseError,
+    Prediction, RenderOptions, Signature, TypeIR,
 };
 
 #[derive(Default, Clone)]
@@ -526,15 +527,19 @@ impl ChatAdapter {
             return String::new();
         };
         let input_output_format = <S::Input as BamlTypeTrait>::baml_output_format();
+        let input_json = build_input_context_value(fields, S::input_fields(), input_output_format);
+        let vars = Value::Object(serde_json::Map::new());
 
         let mut result = String::new();
         for field_spec in S::input_fields() {
             if let Some(value) = fields.get(field_spec.rust_name) {
                 result.push_str(&format!("[[ ## {} ## ]]\n", field_spec.name));
-                result.push_str(&format_baml_value_for_prompt_typed(
+                result.push_str(&render_input_field(
+                    field_spec,
                     value,
+                    &input_json,
                     input_output_format,
-                    field_spec.format,
+                    &vars,
                 ));
                 result.push_str("\n\n");
             }
@@ -880,23 +885,92 @@ fn format_baml_value_for_prompt(value: &BamlValue) -> String {
     }
 }
 
-fn format_baml_value_for_prompt_typed(
+fn render_input_field(
+    field_spec: &crate::FieldSpec,
     value: &BamlValue,
+    input: &Value,
     output_format: &OutputFormatContent,
-    format: Option<&str>,
+    vars: &Value,
 ) -> String {
-    let format = match format {
-        Some(format) => format,
-        None => {
-            if let BamlValue::String(s) = value {
-                return s.clone();
-            }
-            "json"
+    match field_spec.input_render {
+        InputRenderSpec::Default => match value {
+            BamlValue::String(s) => s.clone(),
+            _ => crate::bamltype::internal_baml_jinja::format_baml_value(
+                value,
+                output_format,
+                "json",
+            )
+            .unwrap_or_else(|_| "<error>".to_string()),
+        },
+        InputRenderSpec::Format(format) => {
+            crate::bamltype::internal_baml_jinja::format_baml_value(value, output_format, format)
+                .unwrap_or_else(|_| "<error>".to_string())
+        }
+        InputRenderSpec::Jinja(template) => {
+            render_input_field_jinja(template, field_spec, value, input, output_format, vars)
+                .unwrap_or_else(|_| "<error>".to_string())
         }
+    }
+}
+
+fn render_input_field_jinja(
+    template: &str,
+    field_spec: &crate::FieldSpec,
+    value: &BamlValue,
+    input: &Value,
+    output_format: &OutputFormatContent,
+    vars: &Value,
+) -> Result<String, minijinja::Error> {
+    let mut env = minijinja::Environment::new();
+    env.set_undefined_behavior(UndefinedBehavior::Strict);
+    env.add_template("__input_field__", template)?;
+    let template = env.get_template("__input_field__")?;
+
+    let this = baml_value_to_render_json(value, output_format);
+    let field = json!({
+        "name": field_spec.name,
+        "rust_name": field_spec.rust_name,
+        "type": (field_spec.type_ir)().diagnostic_repr().to_string(),
+    });
+    let context = json!({
+        "this": this,
+        "input": input,
+        "field": field,
+        "vars": vars,
+    });
+
+    template.render(minijinja::Value::from_serialize(context))
+}
+
+fn build_input_context_value(
+    fields: &crate::bamltype::baml_types::BamlMap<String, BamlValue>,
+    field_specs: &[crate::FieldSpec],
+    output_format: &OutputFormatContent,
+) -> Value {
+    let mut map = serde_json::Map::new();
+
+    for field_spec in field_specs {
+        let Some(value) = fields.get(field_spec.rust_name) else {
+            continue;
+        };
+        let value_json = baml_value_to_render_json(value, output_format);
+        map.insert(field_spec.rust_name.to_string(), value_json.clone());
+        if field_spec.name != field_spec.rust_name {
+            map.entry(field_spec.name.to_string()).or_insert(value_json);
+        }
+    }
+
+    Value::Object(map)
+}
+
+fn baml_value_to_render_json(value: &BamlValue, output_format: &OutputFormatContent) -> Value {
+    let Ok(rendered_json) =
+        crate::bamltype::internal_baml_jinja::format_baml_value(value, output_format, "json")
+    else {
+        return serde_json::to_value(value).unwrap_or(Value::Null);
     };
 
-    crate::bamltype::internal_baml_jinja::format_baml_value(value, output_format, format)
-        .unwrap_or_else(|_| "<error>".to_string())
+    serde_json::from_str(&rendered_json).unwrap_or(Value::Null)
 }
 
 fn collect_flags_recursive(value: &BamlValueWithFlags, flags: &mut Vec<Flag>) {

crates/dspy-rs/src/adapter/chat.rs

@@ -2,14 +2,21 @@ use crate::{Example, OutputFormatContent, TypeIR};
 use anyhow::Result;
 use serde_json::Value;
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum InputRenderSpec {
+    Default,
+    Format(&'static str),
+    Jinja(&'static str),
+}
+
 #[derive(Debug, Clone, Copy)]
 pub struct FieldSpec {
     pub name: &'static str,
     pub rust_name: &'static str,
     pub description: &'static str,
     pub type_ir: fn() -> TypeIR,
     pub constraints: &'static [ConstraintSpec],
-    pub format: Option<&'static str>,
+    pub input_render: InputRenderSpec,
 }
 
 #[derive(Debug, Clone, Copy)]

crates/dspy-rs/src/core/signature.rs

@@ -10,7 +10,7 @@ use tracing::{debug, trace};
 use crate::adapter::Adapter;
 use crate::bamltype::baml_types::BamlMap;
 use crate::bamltype::compat::{BamlValueConvert, ToBamlValue};
-use crate::core::{FieldSpec, MetaSignature, Module, Optimizable, Signature};
+use crate::core::{FieldSpec, InputRenderSpec, MetaSignature, Module, Optimizable, Signature};
 use crate::{
     BamlValue, CallResult, Chat, ChatAdapter, Example, GLOBAL_SETTINGS, LM, LmError, LmUsage,
     PredictError, Prediction,
@@ -258,8 +258,14 @@ fn field_specs_to_value(fields: &[FieldSpec], field_type: &'static str) -> Value
         meta.insert("desc".to_string(), json!(field.description));
         meta.insert("schema".to_string(), json!(""));
         meta.insert("__dsrs_field_type".to_string(), json!(field_type));
-        if let Some(format) = field.format {
-            meta.insert("format".to_string(), json!(format));
+        match field.input_render {
+            InputRenderSpec::Default => {}
+            InputRenderSpec::Format(format) => {
+                meta.insert("format".to_string(), json!(format));
+            }
+            InputRenderSpec::Jinja(template) => {
+                meta.insert("render".to_string(), json!({ "jinja": template }));
+            }
         }
         result.insert(field.rust_name.to_string(), Value::Object(meta));
     }

crates/dspy-rs/src/predictors/predict.rs

@@ -62,6 +62,34 @@ struct DefaultFormatSig {
     answer: String,
 }
 
+#[derive(Signature, Clone, Debug)]
+/// Render a context field using Jinja.
+struct RenderJinjaSig {
+    #[input]
+    question: String,
+
+    #[input]
+    #[alias("ctx")]
+    #[render(
+        jinja = "{{ this.text }} | {{ input.question }} | {{ input.ctx.text }} | {{ input.context.text }} | {{ field.name }} | {{ field.rust_name }}"
+    )]
+    context: Document,
+
+    #[output]
+    answer: String,
+}
+
+#[derive(Signature, Clone, Debug)]
+/// Render with strict undefined vars.
+struct RenderJinjaStrictSig {
+    #[input]
+    #[render(jinja = "{{ missing_var }}")]
+    question: String,
+
+    #[output]
+    answer: String,
+}
+
 fn extract_field(message: &str, field_name: &str) -> String {
     let start_marker = format!("[[ ## {field_name} ## ]]");
     let start_pos = message
@@ -184,3 +212,35 @@ fn typed_input_default_non_string_is_json() {
         .expect("expected array with object");
     assert_eq!(first.get("text").and_then(|v| v.as_str()), Some("Hello"));
 }
+
+#[test]
+fn typed_input_render_jinja_uses_context_values() {
+    let adapter = ChatAdapter;
+    let input = RenderJinjaSigInput {
+        question: "Question".to_string(),
+        context: Document {
+            text: "Hello".to_string(),
+        },
+    };
+
+    let message = adapter.format_user_message_typed::<RenderJinjaSig>(&input);
+    let context_value = extract_field(&message, "ctx");
+
+    assert_eq!(
+        context_value,
+        "Hello | Question | Hello | Hello | ctx | context"
+    );
+}
+
+#[test]
+fn typed_input_render_jinja_strict_undefined_returns_error_sentinel() {
+    let adapter = ChatAdapter;
+    let input = RenderJinjaStrictSigInput {
+        question: "Question".to_string(),
+    };
+
+    let message = adapter.format_user_message_typed::<RenderJinjaStrictSig>(&input);
+    let question_value = extract_field(&message, "question");
+
+    assert_eq!(question_value, "<error>");
+}

crates/dspy-rs/tests/test_input_format.rs

@@ -19,6 +19,7 @@ quote = "1"
 proc-macro2 = "1"
 proc-macro-crate = "3.2"
 serde_json = { version = "1.0.143", features = ["preserve_order"] }
+minijinja = { git = "https://github.com/boundaryml/minijinja.git", branch = "main", default-features = false, features = ["serde"] }
 
 [dev-dependencies]
 dspy-rs = { path = "../dspy-rs" }