Skip to content

Commit 4153bd3

Browse files
antoniosarosiantoniosarosi
authored
Maintain attribute information when combining fields (#1585)
Fix #1256 <!-- ELLIPSIS_HIDDEN --> ---- > [!IMPORTANT] > Maintain attribute information in union types by associating attributes from the last type in the union. > > - **Behavior**: > - `combine_field_types` in `parse_field.rs` now retains attributes from the last type in a union. > - Handles union types like `string? | int @alias("hello")` by associating the alias with the union. > - **Tests**: > - Adds `test_class_property_alias` in `test_runtime.rs` to verify alias handling in union types. > - Updates existing tests in `parse_field.rs` to reflect changes in attribute handling. > - **Misc**: > - Minor refactoring in `repr.rs` to simplify attribute extraction logic. > > <sup>This description was created by </sup>[<img alt="Ellipsis" src="https://img.shields.io/badge/Ellipsis-blue?color=175173">](https://www.ellipsis.dev?ref=BoundaryML%2Fbaml&utm_source=github&utm_medium=referral)<sup> for 67a3549. It will automatically update as commits are pushed.</sup> <!-- ELLIPSIS_HIDDEN -->
1 parent 6c6c5c2 commit 4153bd3

4 files changed

Lines changed: 124 additions & 82 deletions

File tree

‎engine/baml-lib/baml-core/src/ir/repr.rs‎

Lines changed: 79 additions & 77 deletions
Original file line numberDiff line numberDiff line change
@@ -453,83 +453,85 @@ fn to_ir_attributes(
453453
db: &ParserDatabase,
454454
maybe_ast_attributes: Option<&Attributes>,
455455
) -> (IndexMap<String, UnresolvedValue<()>>, Vec<Constraint>) {
456-
let null_result = (IndexMap::new(), Vec::new());
457-
maybe_ast_attributes.map_or(null_result, |attributes| {
458-
let Attributes {
459-
description,
460-
alias,
461-
dynamic_type,
462-
skip,
463-
constraints,
464-
streaming_done,
465-
streaming_needed,
466-
streaming_state,
467-
} = attributes;
468-
469-
let description = description
470-
.as_ref()
471-
.map(|d| ("description".to_string(), d.without_meta()));
472-
473-
let alias = alias
474-
.as_ref()
475-
.map(|v| ("alias".to_string(), v.without_meta()));
476-
477-
let dynamic_type = dynamic_type.as_ref().and_then(|v| {
478-
if *v {
479-
Some(("dynamic_type".to_string(), UnresolvedValue::Bool(true, ())))
480-
} else {
481-
None
482-
}
483-
});
484-
let skip = skip.as_ref().and_then(|v| {
485-
if *v {
486-
Some(("skip".to_string(), UnresolvedValue::Bool(true, ())))
487-
} else {
488-
None
489-
}
490-
});
491-
let streaming_done = streaming_done.as_ref().and_then(|v| {
492-
if *v {
493-
Some(("stream.done".to_string(), UnresolvedValue::Bool(true, ())))
494-
} else {
495-
None
496-
}
497-
});
498-
let streaming_needed = streaming_needed.as_ref().and_then(|v| {
499-
if *v {
500-
Some((
501-
"stream.not_null".to_string(),
502-
UnresolvedValue::Bool(true, ()),
503-
))
504-
} else {
505-
None
506-
}
507-
});
508-
let streaming_state = streaming_state.as_ref().and_then(|v| {
509-
if *v {
510-
Some((
511-
"stream.with_state".to_string(),
512-
UnresolvedValue::Bool(true, ()),
513-
))
514-
} else {
515-
None
516-
}
517-
});
518-
519-
let meta = vec![
520-
description,
521-
alias,
522-
dynamic_type,
523-
skip,
524-
streaming_done,
525-
streaming_needed,
526-
streaming_state,
527-
]
528-
.into_iter()
529-
.filter_map(|s| s)
530-
.collect();
531-
(meta, constraints.clone())
532-
})
456+
let Some(attributes) = maybe_ast_attributes else {
457+
return (IndexMap::new(), Vec::new());
458+
};
459+
460+
let Attributes {
461+
description,
462+
alias,
463+
dynamic_type,
464+
skip,
465+
constraints,
466+
streaming_done,
467+
streaming_needed,
468+
streaming_state,
469+
} = attributes;
470+
471+
let description = description
472+
.as_ref()
473+
.map(|d| ("description".to_string(), d.without_meta()));
474+
475+
let alias = alias
476+
.as_ref()
477+
.map(|v| ("alias".to_string(), v.without_meta()));
478+
479+
let dynamic_type = dynamic_type.as_ref().and_then(|v| {
480+
if *v {
481+
Some(("dynamic_type".to_string(), UnresolvedValue::Bool(true, ())))
482+
} else {
483+
None
484+
}
485+
});
486+
let skip = skip.as_ref().and_then(|v| {
487+
if *v {
488+
Some(("skip".to_string(), UnresolvedValue::Bool(true, ())))
489+
} else {
490+
None
491+
}
492+
});
493+
let streaming_done = streaming_done.as_ref().and_then(|v| {
494+
if *v {
495+
Some(("stream.done".to_string(), UnresolvedValue::Bool(true, ())))
496+
} else {
497+
None
498+
}
499+
});
500+
let streaming_needed = streaming_needed.as_ref().and_then(|v| {
501+
if *v {
502+
Some((
503+
"stream.not_null".to_string(),
504+
UnresolvedValue::Bool(true, ()),
505+
))
506+
} else {
507+
None
508+
}
509+
});
510+
let streaming_state = streaming_state.as_ref().and_then(|v| {
511+
if *v {
512+
Some((
513+
"stream.with_state".to_string(),
514+
UnresolvedValue::Bool(true, ()),
515+
))
516+
} else {
517+
None
518+
}
519+
});
520+
521+
let meta = vec![
522+
description,
523+
alias,
524+
dynamic_type,
525+
skip,
526+
streaming_done,
527+
streaming_needed,
528+
streaming_state,
529+
]
530+
.into_iter()
531+
.filter_map(|s| s)
532+
.collect();
533+
534+
(meta, constraints.clone())
533535
}
534536

535537
/// Nodes allow attaching metadata to a given IR entity: attributes, source location, etc

‎engine/baml-lib/schema-ast/src/parser/parse_field.rs‎

Lines changed: 18 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,7 @@ pub(crate) fn parse_value_expr(
5454
_ => Err(DatamodelError::new_model_validation_error(
5555
"This field declaration is invalid. It is either missing a name or a type.",
5656
container_type,
57-
model_name.as_ref().map_or("<unknown>", |f| f.name()),
57+
model_name.as_ref().map_or("<unknown>", Identifier::name),
5858
diagnostics.span(pair_span),
5959
)),
6060
}
@@ -219,13 +219,21 @@ fn combine_field_types(types: Vec<FieldType>) -> Option<FieldType> {
219219

220220
let mut seen_types = vec![combined_type.clone()];
221221

222+
// In a union, use the attributes associated with the last type as the
223+
// attributes of the union. Example:
224+
//
225+
// field: string? | int @alias("hello")
226+
//
227+
// The alias is part of the union.
228+
let last_field_attrs = types.last().map(|t| t.attributes().to_vec());
229+
222230
let mut earliest_start = combined_type.span().start;
223231
let mut latest_end = combined_type.span().end;
224232

225233
for next_type in types.into_iter().skip(1) {
226-
seen_types.push(next_type.clone());
234+
let span = next_type.span().to_owned();
235+
seen_types.push(next_type);
227236

228-
let span = next_type.span();
229237
if span.start < earliest_start {
230238
earliest_start = span.start;
231239
}
@@ -245,6 +253,11 @@ fn combine_field_types(types: Vec<FieldType>) -> Option<FieldType> {
245253
);
246254
}
247255

256+
// We know it's a union because it was assigned above in the for loop.
257+
if let FieldType::Union(_, _, _, attrs) = &mut combined_type {
258+
*attrs = last_field_attrs;
259+
}
260+
248261
Some(combined_type)
249262
}
250263

@@ -531,9 +544,9 @@ mod tests {
531544
name: (value, Span::fake()).into(),
532545
parenthesized: false,
533546
arguments: ArgumentsList {
534-
arguments: Vec::new()
547+
arguments: Vec::new(),
535548
},
536-
span: Span::fake()
549+
span: Span::fake(),
537550
}
538551
}
539552
}

‎engine/baml-runtime/src/internal/prompt_renderer/render_output_format.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ pub fn render_output_format(
2929
.build())
3030
}
3131

32+
#[derive(Debug)]
3233
enum OverridableValue<T> {
3334
Unset,
3435
SetEmpty,

‎engine/baml-runtime/tests/test_runtime.rs‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1007,4 +1007,30 @@ test RecursiveAliasCycle {
10071007
"##,
10081008
})
10091009
}
1010+
1011+
#[test]
1012+
fn test_class_property_alias() -> anyhow::Result<()> {
1013+
run_type_builder_block_test(TypeBuilderBlockTest {
1014+
function_name: "Fn",
1015+
test_name: "Test",
1016+
baml: r##"
1017+
class PropertyAlias {
1018+
property string? | int @alias("hello")
1019+
}
1020+
1021+
function Fn() -> PropertyAlias {
1022+
client "openai/gpt-4o"
1023+
prompt #"
1024+
{{ctx.output_format}}
1025+
"#
1026+
}
1027+
1028+
test Test {
1029+
functions [Fn]
1030+
args {
1031+
}
1032+
}
1033+
"##,
1034+
})
1035+
}
10101036
}

0 commit comments

Comments
 (0)