diff --git a/engine/baml-lib/jinja-runtime/src/lib.rs b/engine/baml-lib/jinja-runtime/src/lib.rs index 75feb75cce..442dd6cd4f 100644 --- a/engine/baml-lib/jinja-runtime/src/lib.rs +++ b/engine/baml-lib/jinja-runtime/src/lib.rs @@ -252,33 +252,75 @@ fn render_minijinja(params: MinijinjaRenderParams) -> Result(media_data) { - Ok(m) => Some(ChatMessagePart::Media(m)), + // Be tolerant of surrounding punctuation/whitespace around media markers + let start_tag = ":baml-start-media:"; + let end_tag = ":baml-end-media:"; + let mut consumed_any = false; + + let mut remaining = part; + loop { + let start_idx = match remaining.find(start_tag) { + Some(i) => i, + None => break, + }; + + let end_search_start = start_idx + start_tag.len(); + let end_idx_rel = match remaining[end_search_start..].find(end_tag) { + Some(i) => i, + None => break, + }; + let end_idx = end_search_start + end_idx_rel; + + // Preceding text + let before = &remaining[..start_idx]; + if !before.trim().is_empty() { + let txt = ChatMessagePart::Text(before.trim().to_string()); + parts.push(match &meta { + Some(m) => txt.with_meta(m.clone()), + None => txt, + }); + } + + // Media JSON between tags + let media_json = &remaining[end_search_start..end_idx]; + match serde_json::from_str::(media_json) { + Ok(m) => { + let media_part = ChatMessagePart::Media(m); + parts.push(match &meta { + Some(m) => media_part.with_meta(m.clone()), + None => media_part, + }); + } Err(_) => Err(minijinja::Error::new( ErrorKind::CannotUnpack, - format!("Media variable had unrecognizable data: {media_data}"), + format!("Media variable had unrecognizable data: {media_json}"), ))?, } - } else if !part.trim().is_empty() { - Some(ChatMessagePart::Text(part.trim().to_string())) + + consumed_any = true; + + // Advance remaining after end tag + let after_start = end_idx + end_tag.len(); + remaining = &remaining[after_start..]; + } + + if consumed_any { + // Whatever remains (tail) after last media extraction + if !remaining.trim().is_empty() { + let txt = ChatMessagePart::Text(remaining.trim().to_string()); + parts.push(match &meta { + Some(m) => txt.with_meta(m.clone()), + None => txt, + }); + } } else { - None - }; - - if let Some(part) = part { - if let Some(meta) = &meta { - parts.push(part.with_meta(meta.clone())); - } else { - parts.push(part); + // No media markers in this segment; treat as text if non-empty + if !part.trim().is_empty() { + let txt = ChatMessagePart::Text(part.trim().to_string()); + parts.push(match &meta { + Some(m) => txt.with_meta(m.clone()), + None => txt, + }); } } } @@ -675,6 +717,66 @@ mod render_tests { Ok(()) } + #[test] + fn render_image_nested_without_outer_delimiter() -> anyhow::Result<()> { + setup_logging(); + let ir = make_test_ir( + " + class C { + + } + ", + )?; + + let args: BamlValue = BamlValue::Map(BamlMap::from([( + "obj".to_string(), + BamlValue::Class( + "SomeClass".to_string(), + BamlMap::from([( + "img".to_string(), + BamlValue::Media(BamlMedia::url( + BamlMediaType::Image, + "https://example.com/image.jpg".to_string(), + None, + )), + )]), + ), + )])); + + // Render the media inside a class render where punctuation can surround the media markers + let rendered = render_prompt( + "{{ _.chat(\"system\") }}\nHere: {{ obj }}", + &args, + RenderContext { + client: RenderContext_Client { + name: "gpt4".to_string(), + provider: "openai".to_string(), + default_role: "system".to_string(), + allowed_roles: vec!["system".to_string()], + remap_role: HashMap::new(), + options: IndexMap::new(), + }, + output_format: OutputFormatContent::new_string(), + tags: HashMap::from([("ROLE".to_string(), BamlValue::String("john doe".into()))]), + }, + &[], + &ir, + &HashMap::new(), + )?; + + match rendered { + RenderedPrompt::Chat(messages) => { + assert_eq!(messages.len(), 1); + let parts = &messages[0].parts; + // Expect at least one media part present even when surrounded by class text + assert!(parts.iter().any(|p| matches!(p, ChatMessagePart::Media(_)))); + } + _ => anyhow::bail!("Expected Chat prompt"), + } + + Ok(()) + } + #[test] fn render_image_suffix() -> anyhow::Result<()> { setup_logging(); diff --git a/integ-tests/go/baml_client/functions.go b/integ-tests/go/baml_client/functions.go index a3e04704a2..26f546ee98 100644 --- a/integ-tests/go/baml_client/functions.go +++ b/integ-tests/go/baml_client/functions.go @@ -15,7 +15,6 @@ package baml_client import ( "context" - "fmt" "example.com/integ-tests/baml_client/types" baml "github.com/boundaryml/baml/engine/language_client_go/pkg" diff --git a/integ-tests/go/baml_client/functions_parse.go b/integ-tests/go/baml_client/functions_parse.go index 3092a05515..099a0f20b2 100644 --- a/integ-tests/go/baml_client/functions_parse.go +++ b/integ-tests/go/baml_client/functions_parse.go @@ -15,8 +15,8 @@ package baml_client import ( "context" - "fmt" + "example.com/integ-tests/baml_client/stream_types" "example.com/integ-tests/baml_client/types" baml "github.com/boundaryml/baml/engine/language_client_go/pkg" ) diff --git a/integ-tests/go/baml_client/functions_parse_stream.go b/integ-tests/go/baml_client/functions_parse_stream.go index 2446714547..a95b15112c 100644 --- a/integ-tests/go/baml_client/functions_parse_stream.go +++ b/integ-tests/go/baml_client/functions_parse_stream.go @@ -15,7 +15,6 @@ package baml_client import ( "context" - "fmt" "example.com/integ-tests/baml_client/stream_types" "example.com/integ-tests/baml_client/types" diff --git a/integ-tests/go/baml_client/functions_stream.go b/integ-tests/go/baml_client/functions_stream.go index de94ea8786..7dcd44cd44 100644 --- a/integ-tests/go/baml_client/functions_stream.go +++ b/integ-tests/go/baml_client/functions_stream.go @@ -15,7 +15,6 @@ package baml_client import ( "context" - "fmt" "example.com/integ-tests/baml_client/stream_types" "example.com/integ-tests/baml_client/types" diff --git a/integ-tests/go/baml_client/runtime.go b/integ-tests/go/baml_client/runtime.go index 41c4bd011b..de6bd35da4 100644 --- a/integ-tests/go/baml_client/runtime.go +++ b/integ-tests/go/baml_client/runtime.go @@ -27,11 +27,9 @@ package baml_client import ( - "fmt" "os" "strings" - "example.com/integ-tests/baml_client/type_builder" baml "github.com/boundaryml/baml/engine/language_client_go/pkg" ) diff --git a/integ-tests/go/baml_client/stream_types/classes.go b/integ-tests/go/baml_client/stream_types/classes.go index f5e6c5d2ad..bd98fe8ec4 100644 --- a/integ-tests/go/baml_client/stream_types/classes.go +++ b/integ-tests/go/baml_client/stream_types/classes.go @@ -14,6 +14,7 @@ package stream_types import ( + "encoding/json" "fmt" baml "github.com/boundaryml/baml/engine/language_client_go/pkg" diff --git a/integ-tests/go/baml_client/stream_types/type_aliases.go b/integ-tests/go/baml_client/stream_types/type_aliases.go index 5b53d24391..da763ee586 100644 --- a/integ-tests/go/baml_client/stream_types/type_aliases.go +++ b/integ-tests/go/baml_client/stream_types/type_aliases.go @@ -14,6 +14,12 @@ package stream_types import ( + "encoding/json" + "fmt" + + baml "github.com/boundaryml/baml/engine/language_client_go/pkg" + "github.com/boundaryml/baml/engine/language_client_go/pkg/cffi" + "example.com/integ-tests/baml_client/types" ) diff --git a/integ-tests/go/baml_client/stream_types/unions.go b/integ-tests/go/baml_client/stream_types/unions.go index 680106a14c..522b6fd0fd 100644 --- a/integ-tests/go/baml_client/stream_types/unions.go +++ b/integ-tests/go/baml_client/stream_types/unions.go @@ -19,6 +19,8 @@ import ( baml "github.com/boundaryml/baml/engine/language_client_go/pkg" "github.com/boundaryml/baml/engine/language_client_go/pkg/cffi" + + "example.com/integ-tests/baml_client/types" ) type Union2EarthlingOrMartian struct { diff --git a/integ-tests/go/baml_client/stream_types/utils.go b/integ-tests/go/baml_client/stream_types/utils.go index 5ab412b0d7..3077ce5978 100644 --- a/integ-tests/go/baml_client/stream_types/utils.go +++ b/integ-tests/go/baml_client/stream_types/utils.go @@ -12,3 +12,11 @@ // $ go install github.com/boundaryml/baml/baml-cli package stream_types + +import ( + "encoding/json" + "fmt" + + baml "github.com/boundaryml/baml/engine/language_client_go/pkg" + "github.com/boundaryml/baml/engine/language_client_go/pkg/cffi" +) diff --git a/integ-tests/go/baml_client/type_builder/type_builder.go b/integ-tests/go/baml_client/type_builder/type_builder.go index 9796eec1c3..1bb14d223f 100644 --- a/integ-tests/go/baml_client/type_builder/type_builder.go +++ b/integ-tests/go/baml_client/type_builder/type_builder.go @@ -13,11 +13,7 @@ package type_builder -import ( - "fmt" - - baml "github.com/boundaryml/baml/engine/language_client_go/pkg" -) +import baml "github.com/boundaryml/baml/engine/language_client_go/pkg" type Type = baml.Type diff --git a/integ-tests/go/baml_client/type_map.go b/integ-tests/go/baml_client/type_map.go index 17362143bc..8f980920f0 100644 --- a/integ-tests/go/baml_client/type_map.go +++ b/integ-tests/go/baml_client/type_map.go @@ -14,8 +14,6 @@ package baml_client import ( - "reflect" - "example.com/integ-tests/baml_client/stream_types" "example.com/integ-tests/baml_client/types" ) diff --git a/integ-tests/go/baml_client/types/classes.go b/integ-tests/go/baml_client/types/classes.go index c8ec9c795e..526fe7c5ac 100644 --- a/integ-tests/go/baml_client/types/classes.go +++ b/integ-tests/go/baml_client/types/classes.go @@ -14,6 +14,7 @@ package types import ( + "encoding/json" "fmt" baml "github.com/boundaryml/baml/engine/language_client_go/pkg" diff --git a/integ-tests/go/baml_client/types/type_aliases.go b/integ-tests/go/baml_client/types/type_aliases.go index 81387dc207..8d2b888956 100644 --- a/integ-tests/go/baml_client/types/type_aliases.go +++ b/integ-tests/go/baml_client/types/type_aliases.go @@ -13,6 +13,14 @@ package types +import ( + "encoding/json" + "fmt" + + baml "github.com/boundaryml/baml/engine/language_client_go/pkg" + "github.com/boundaryml/baml/engine/language_client_go/pkg/cffi" +) + type Amount = int64 type Combination = Union6BoolOrFloatOrIntOrListStringOrMapStringKeyListStringValueOrString