Skip to content

Commit dfd8a85

Browse files
authored
Merge pull request #79 from krypticmouse/codex/input-render-jinja-support
[codex] add input-field jinja rendering support
2 parents aad5952 + 0e3f648 commit dfd8a85

33 files changed

Lines changed: 770 additions & 48 deletions

‎Cargo.lock‎

Lines changed: 2 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎const_block_generic‎

-453 KB
Binary file not shown.

‎crates/bamltype/tests/ui.rs‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,8 @@
11
#[test]
2+
#[cfg_attr(
3+
miri,
4+
ignore = "trybuild launches subprocesses and is unsupported under miri"
5+
)]
26
fn ui_compile_failures() {
37
let t = trybuild::TestCases::new();
48
t.compile_fail("tests/ui/*.rs");

‎crates/dspy-rs/Cargo.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@ rig-core = { git = "https://github.com/0xPlaygrounds/rig", rev="e7849df" }
4444
enum_dispatch = "0.3.13"
4545
tracing = "0.1.44"
4646
tracing-subscriber = { version = "0.3.22", features = ["env-filter", "fmt"] }
47+
minijinja = { git = "https://github.com/boundaryml/minijinja.git", branch = "main", default-features = false, features = ["builtins", "serde"] }
4748

4849
[package.metadata.cargo-machete]
4950
ignored = ["rig-core"]

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

Lines changed: 214 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -4,15 +4,20 @@ use bamltype::jsonish::BamlValueWithFlags;
44
use bamltype::jsonish::deserializer::coercer::run_user_checks;
55
use bamltype::jsonish::deserializer::deserialize_flags::DeserializerConditions;
66
use indexmap::IndexMap;
7+
use minijinja::UndefinedBehavior;
8+
use minijinja::value::{Kwargs, Value as MiniJinjaValue};
79
use regex::Regex;
8-
use std::sync::LazyLock;
10+
use serde_json::{Value, json};
11+
use std::collections::HashMap;
12+
use std::sync::{LazyLock, Mutex};
913
use tracing::{debug, trace};
1014

1115
use super::Adapter;
1216
use crate::CallMetadata;
1317
use 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;
3338
static 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+
36136
fn 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

8871085
fn collect_flags_recursive(value: &BamlValueWithFlags, flags: &mut Vec<Flag>) {

‎crates/dspy-rs/src/core/mod.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,7 @@ pub use lm::*;
3737
pub use module::*;
3838
pub use module_ext::*;
3939
pub use predicted::{CallMetadata, ConstraintResult, FieldMeta, Predicted};
40-
pub use schema::{FieldMetadataSpec, FieldPath, FieldSchema, SignatureSchema};
40+
pub use schema::{FieldMetadataSpec, FieldPath, FieldSchema, InputRenderSpec, SignatureSchema};
4141
pub use settings::*;
4242
pub use signature::*;
4343
pub use specials::*;

‎crates/dspy-rs/src/core/schema.rs‎

Lines changed: 19 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,17 @@ impl FieldPath {
4848
/// Static metadata for a single signature field, emitted by `#[derive(Signature)]`.
4949
///
5050
/// Carries the Rust field name, optional LM-facing alias, constraint specs, and
51-
/// format hints. Fed into [`SignatureSchema`] construction alongside Facet shape data.
51+
/// input render hints. Fed into [`SignatureSchema`] construction alongside Facet shape data.
52+
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
53+
pub enum InputRenderSpec {
54+
/// Default behavior: strings are raw, non-strings are rendered as JSON.
55+
Default,
56+
/// Explicit format hint (`#[format("json" | "yaml" | "toon")]`).
57+
Format(&'static str),
58+
/// Custom Jinja template (`#[render(jinja = "...")]`).
59+
Jinja(&'static str),
60+
}
61+
5262
#[derive(Debug, Clone, Copy)]
5363
pub struct FieldMetadataSpec {
5464
/// The Rust field name as written in the signature struct.
@@ -57,8 +67,8 @@ pub struct FieldMetadataSpec {
5767
pub alias: Option<&'static str>,
5868
/// Constraint specs from `#[check(...)]` and `#[assert(...)]` attributes.
5969
pub constraints: &'static [ConstraintSpec],
60-
/// Optional format hint (e.g. `#[format = "json"]`).
61-
pub format: Option<&'static str>,
70+
/// Input rendering policy for this field.
71+
pub input_render: InputRenderSpec,
6272
}
6373

6474
/// Complete schema for a single field in a signature, combining Facet shape data with metadata.
@@ -81,8 +91,8 @@ pub struct FieldSchema {
8191
pub path: FieldPath,
8292
/// Constraints declared on this field.
8393
pub constraints: &'static [ConstraintSpec],
84-
/// Optional format hint.
85-
pub format: Option<&'static str>,
94+
/// Input rendering policy.
95+
pub input_render: InputRenderSpec,
8696
}
8797

8898
impl FieldSchema {
@@ -329,7 +339,9 @@ fn emit_field(
329339
let lm_name = inherited
330340
.and_then(|meta| meta.alias)
331341
.unwrap_or_else(|| field.effective_name());
332-
let format = inherited.and_then(|meta| meta.format);
342+
let input_render = inherited
343+
.map(|meta| meta.input_render)
344+
.unwrap_or(InputRenderSpec::Default);
333345

334346
out.push(FieldSchema {
335347
lm_name,
@@ -339,7 +351,7 @@ fn emit_field(
339351
shape: field.shape(),
340352
path,
341353
constraints,
342-
format,
354+
input_render,
343355
});
344356

345357
Ok(())

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -73,9 +73,9 @@ pub trait Signature: Send + Sync + 'static {
7373
/// The Facet shape of the output struct.
7474
fn output_shape() -> &'static Shape;
7575

76-
/// Per-field metadata for input fields (aliases, constraints, format hints).
76+
/// Per-field metadata for input fields (aliases, constraints, input render hints).
7777
fn input_field_metadata() -> &'static [FieldMetadataSpec];
78-
/// Per-field metadata for output fields (aliases, constraints, format hints).
78+
/// Per-field metadata for output fields (aliases, constraints, input render hints).
7979
fn output_field_metadata() -> &'static [FieldMetadataSpec];
8080

8181
/// The output format descriptor used by the adapter for structured output parsing.

0 commit comments

Comments
 (0)