diff --git a/baml_language/Cargo.lock b/baml_language/Cargo.lock index 242d6fc41b1..e730fbb4cd7 100644 --- a/baml_language/Cargo.lock +++ b/baml_language/Cargo.lock @@ -1476,7 +1476,6 @@ name = "bex_sap" version = "0.0.0-beta" dependencies = [ "baml_type", - "bex_external_types", "bex_str", "bex_vm_types", "derive_more", @@ -1516,6 +1515,7 @@ dependencies = [ "baml_type", "base64", "bex_heap", + "bex_sap", "bex_str", "bex_vm_types", "btel_clock", @@ -7424,7 +7424,6 @@ dependencies = [ "bex_external_types", "bex_heap", "bex_resource_types", - "bex_sap", "bex_vm_types", "indexmap", "num-bigint", diff --git a/baml_language/crates/baml_builtins2/baml_std/baml/ns_sap/sap.baml b/baml_language/crates/baml_builtins2/baml_std/baml/ns_sap/sap.baml index 31c3375bf60..0c437523e66 100644 --- a/baml_language/crates/baml_builtins2/baml_std/baml/ns_sap/sap.baml +++ b/baml_language/crates/baml_builtins2/baml_std/baml/ns_sap/sap.baml @@ -8,14 +8,16 @@ class _ParseCache { _data: $rust_type, + //baml:mut_vm /// (internal) Final SAP operation shared by streams and public SAP. function _parse_final(self, json: string) -> T throws root.errors.LlmClient { - $rust_io_function + $rust_function } + //baml:mut_vm /// (internal) Partial SAP operation used by `ai.stream.Stream`. function _parse_partial(self, json: string) -> T | _NoYield throws root.errors.LlmClient { - $rust_io_function + $rust_function } } @@ -23,6 +25,8 @@ class _ParseCache { class _NoYield { } +//baml:mut_vm +//baml:fallible /// (internal) Builds a parse cache for `T`. /// /// A free function rather than a `_ParseCache` static method: a *static* @@ -31,7 +35,7 @@ class _NoYield { /// operand slots the call site never fills. Instance methods on a generic /// class (`_parse_final`, `_parse_partial`) are threaded correctly. function _new_parse_cache() -> _ParseCache throws never { - $rust_io_function + $rust_function } /// Parse an LLM's raw text output into `T` with Schema-Aligned Parsing. diff --git a/baml_language/crates/baml_builtins2_codegen/src/extract.rs b/baml_language/crates/baml_builtins2_codegen/src/extract.rs index 08723fb013a..43425fffe30 100644 --- a/baml_language/crates/baml_builtins2_codegen/src/extract.rs +++ b/baml_language/crates/baml_builtins2_codegen/src/extract.rs @@ -1579,8 +1579,15 @@ implements Modifier for Root { .find(|b| b.path == "baml.http._fetch") .unwrap(); assert_eq!(http_fetch.throws, throws(&["Io", "Timeout"])); + } + + /// Schema-aligned parsing does no IO, so its entry points are VM natives. + #[test] + fn test_sap_parse_is_a_vm_builtin() { + let (vm_builtins, io_builtins, _class_defs) = extract_native_builtins().unwrap(); - let sap_final = io_builtins + assert!(!io_builtins.iter().any(|b| b.path.starts_with("baml.sap."))); + let sap_final = vm_builtins .iter() .find(|b| b.path == "baml.sap._ParseCache._parse_final") .unwrap(); diff --git a/baml_language/crates/baml_cli/src/snapshots/baml_cli__describe_command_tests__render_builtin_package_listing.snap b/baml_language/crates/baml_cli/src/snapshots/baml_cli__describe_command_tests__render_builtin_package_listing.snap index f8e38098e22..0105f657791 100644 --- a/baml_language/crates/baml_cli/src/snapshots/baml_cli__describe_command_tests__render_builtin_package_listing.snap +++ b/baml_language/crates/baml_cli/src/snapshots/baml_cli__describe_command_tests__render_builtin_package_listing.snap @@ -216,8 +216,8 @@ class baml.regex.Error /baml/ns_regex/regex. class baml.regex.Regex /baml/ns_regex/regex.baml:97 function baml.regex.new /baml/ns_regex/regex.baml:244 function baml.regex.escape /baml/ns_regex/regex.baml:268 -function baml.sap.parse /baml/ns_sap/sap.baml:41 -function baml.sap.parse_type /baml/ns_sap/sap.baml:54 +function baml.sap.parse /baml/ns_sap/sap.baml:45 +function baml.sap.parse_type /baml/ns_sap/sap.baml:58 class baml.spawn.CancelToken /baml/ns_spawn/spawn.baml:23 interface baml.spawn.Modifier /baml/ns_spawn/spawn.baml:60 class baml.spawn.Plan /baml/ns_spawn/spawn.baml:78 diff --git a/baml_language/crates/baml_compiler2_mir/src/lower.rs b/baml_language/crates/baml_compiler2_mir/src/lower.rs index 535a9d23390..f15dbba8250 100644 --- a/baml_language/crates/baml_compiler2_mir/src/lower.rs +++ b/baml_language/crates/baml_compiler2_mir/src/lower.rs @@ -805,6 +805,37 @@ pub(crate) fn ty_to_pattern_template_from_resolved_ty(ty: &RuntimeTy) -> Option< } } +/// The value an enclosing `match` scrutinizes, for the coarse-tag proofs. +struct MatchScrutinee { + /// The local the match evaluated its scrutinee into, then each snapshot a + /// narrowing arm (`let x: T`) copies it into and tests in its place. + locals: Vec, + /// The scrutinee's resolved static type. + ty: RuntimeTy, +} + +impl MatchScrutinee { + fn new(local: Local, ty: RuntimeTy) -> Self { + Self { + locals: vec![local], + ty, + } + } + + /// The scrutinee's static type when `local` holds the scrutinee. + fn ty_held_by(&self, local: Local) -> Option<&RuntimeTy> { + self.locals.contains(&local).then_some(&self.ty) + } + + /// Records that `copy` was assigned a copy of `source`: if `source` holds + /// the scrutinee, so does `copy`. + fn record_copy(&mut self, source: Local, copy: Local) { + if self.locals.contains(&source) { + self.locals.push(copy); + } + } +} + /// Flatten a `RuntimeTy` union into its leaf members (recursively), pushing each /// non-union leaf into `out`. A recursive alias is kept as a single `TypeAlias` /// leaf — it is not expanded — matching `ResolvedAliases::convert`. @@ -1693,14 +1724,14 @@ struct LoweringContext<'db> { // baml.ops.Compare`). Saved/restored across nested lambdas. lambda_param_tir_types: FxHashMap, - /// The `(local, resolved static type)` of the value the enclosing `match` is - /// scrutinizing, set for the duration of its arm lowering (save/restored - /// across nested matches). A container arm's runtime type test consults this - /// — guarded on the local matching — to decide whether a coarse `LIST`/`MAP` - /// tag suffices in place of a structural element check (see + /// The value the enclosing `match` is scrutinizing, set for the duration of + /// its arm lowering (save/restored across nested matches). A container + /// arm's runtime type test consults this — guarded on the tested local + /// holding the scrutinee — to decide whether a coarse `LIST`/`MAP` tag + /// suffices in place of a structural element check (see /// [`parametric_arm_tag_sufficient`]). `None` outside a match, and the local /// guard keeps `is` / standalone tests element-precise even inside one. - match_scrutinee: Option<(Local, RuntimeTy)>, + match_scrutinee: Option, /// Stable values read while testing patterns, keyed by pattern identity. /// Binding lowering reuses these exact values instead of rereading mutable @@ -12507,9 +12538,10 @@ impl<'db> LoweringContext<'db> { // Expose the scrutinee's static type to each arm's container type test so // it can decide whether a coarse `LIST`/`MAP` tag suffices; restore the // enclosing match's scrutinee (if any) once the arms are lowered. - let saved_scrutinee = self - .match_scrutinee - .replace((scrutinee_local, self.expr_ty(scrutinee))); + let saved_scrutinee = self.match_scrutinee.replace(MatchScrutinee::new( + scrutinee_local, + self.expr_ty(scrutinee), + )); let switched = self.try_lower_as_switch( scrutinee_local, @@ -12585,8 +12617,8 @@ impl<'db> LoweringContext<'db> { let scrutinee_static_ty: Option = self .match_scrutinee .as_ref() - .filter(|(local, _)| *local == scrutinee) - .map(|(_, ty)| ty.clone()); + .and_then(|match_scrutinee| match_scrutinee.ty_held_by(scrutinee)) + .cloned(); // Classify arms: collect (key, arm_index) for int literal, enum // variant, or type patterns, and check for a trailing wildcard/binding. @@ -13208,14 +13240,14 @@ impl<'db> LoweringContext<'db> { // If this is a parametric arm whose type arguments the emitter would // otherwise check structurally, but the enclosing match's scrutinee // statically proves a coarse `LIST`/`MAP`/`FUTURE` tag suffices, - // emit the tag test directly. Guarded on the scrutinee local so an + // emit the tag test directly. Guarded on the tested local holding + // the scrutinee (itself, or a narrowing arm's snapshot of it) so an // `is` / nested test against a different value stays arg-precise. if self .match_scrutinee .as_ref() - .is_some_and(|(local, scrutinee_ty)| { - *local == scrutinee && parametric_arm_tag_sufficient(&ty, scrutinee_ty) - }) + .and_then(|match_scrutinee| match_scrutinee.ty_held_by(scrutinee)) + .is_some_and(|scrutinee_ty| parametric_arm_tag_sufficient(&ty, scrutinee_ty)) { let tag = match &ty { RuntimeTy::List(..) => baml_type::typetag::LIST, @@ -13770,6 +13802,9 @@ impl<'db> LoweringContext<'db> { Place::local(snapshot), Rvalue::Use(Operand::Copy(Place::Local(scrutinee))), ); + if let Some(match_scrutinee) = self.match_scrutinee.as_mut() { + match_scrutinee.record_copy(scrutinee, snapshot); + } self.tested_pattern_values .insert(self.pat_metadata_key(pat_id), snapshot); snapshot diff --git a/baml_language/crates/baml_tests/baml_src/ns_json_alias/json_alias.baml b/baml_language/crates/baml_tests/baml_src/ns_json_alias/json_alias.baml index 5144d1e6726..64b5f55a508 100644 --- a/baml_language/crates/baml_tests/baml_src/ns_json_alias/json_alias.baml +++ b/baml_language/crates/baml_tests/baml_src/ns_json_alias/json_alias.baml @@ -137,3 +137,70 @@ function json_wrap_tag() -> string { test "impl_on_recursive_alias_generic_arg_dispatches" { assert.equal(json_wrap_tag(), "json-wrap") } + +// ─── 1.9 Container arms on a json scrutinee ────────────────────────────────── + +// `json[]` and `map` are json's only list and map members, so a +// match on a json value tests these arms by the runtime tag alone. Every kind +// of json value must still reach its own arm, also through recursion (the +// request builders prune nulls this way). + +function json_kind(j: json) -> string { + match (j) { + let object: map => "map", + let items: json[] => "list", + _ => "scalar", + } +} + +function json_kinds() -> string[] { + [ + json_kind({"a": 1}), + json_kind({}), + json_kind([1, "two"]), + json_kind([]), + json_kind("map"), + json_kind(1), + json_kind(true), + json_kind(null), + ] +} + +test "json_container_arms_match_each_kind" { + assert.equal(json_kinds(), ["map", "map", "list", "list", "scalar", "scalar", "scalar", "scalar"]) +} + +function drop_nulls(j: json) -> json { + match (j) { + let object: map => { + let out: map = {}; + for (let key in object.keys()) { + let entry = object.get(key); + match (entry) { + null => null, + _ => { + let _ = out.set(key, drop_nulls(entry)); + null + }, + }; + } + out + }, + let items: json[] => { + let out: json[] = []; + for (let item in items) { + out.push(drop_nulls(item)); + } + out + }, + _ => j, + } +} + +function drop_nulls_nested() -> string { + baml.json.stringify(drop_nulls(baml.json.parse("{\"a\":null,\"b\":[1,{\"c\":null,\"d\":\"x\"}],\"e\":{\"f\":null}}"))) +} + +test "json_container_arms_recurse_through_nested_values" { + assert.equal(drop_nulls_nested(), "{\"b\":[1,{\"d\":\"x\"}],\"e\":{}}") +} diff --git a/baml_language/crates/baml_tests/baml_src/ns_parse_companions/parse_companions.baml b/baml_language/crates/baml_tests/baml_src/ns_parse_companions/parse_companions.baml index fad5822eeb1..133789c83c1 100644 --- a/baml_language/crates/baml_tests/baml_src/ns_parse_companions/parse_companions.baml +++ b/baml_language/crates/baml_tests/baml_src/ns_parse_companions/parse_companions.baml @@ -172,3 +172,49 @@ test "sap_parse_matches_the_enum_lines_the_prompt_renders" { assert.is_true(prompt.includes("- A: eaten in the morning")); assert.is_true(prompt.includes("- B: LUNCH eaten at noon")) } + +// ─── @skip fields ──────────────────────────────────────────────────────────── +// +// The parser never reads a `@skip` field (the prompt does not show it), so a +// parsed instance holds its type's empty value: "" for a string, the zero of a +// number, false, an empty list, null for an optional field, an enum's first +// variant, and a class of its fields' empty values. + +enum ParseCompanionsSkipColor { + Red, + Green, +} + +class ParseCompanionsSkipInner { + label: string, + count: int, +} + +class ParseCompanionsSkipPerson { + name: string, + age: int @skip, + score: float @skip, + active: bool @skip, + tags: string[] @skip, + nick: string? @skip, + color: ParseCompanionsSkipColor @skip, + inner: ParseCompanionsSkipInner @skip, +} + +function parse_companions_skipped_fields_fn() -> bool { + let parsed = baml.sap.parse("{\"name\": \"Ada\", \"age\": 37}"); + parsed == ParseCompanionsSkipPerson { + name: "Ada", + age: 0, + score: 0.0, + active: false, + tags: [], + nick: null, + color: ParseCompanionsSkipColor.Red, + inner: ParseCompanionsSkipInner { label: "", count: 0 }, + } +} + +test "sap_parse_fills_skipped_fields_with_empty_values" { + assert.is_true(parse_companions_skipped_fields_fn()) +} diff --git a/baml_language/crates/baml_tests/src/compiler2_mir/mod.rs b/baml_language/crates/baml_tests/src/compiler2_mir/mod.rs index e8d7dc6cbcb..3740955607f 100644 --- a/baml_language/crates/baml_tests/src/compiler2_mir/mod.rs +++ b/baml_language/crates/baml_tests/src/compiler2_mir/mod.rs @@ -23,6 +23,25 @@ fn make_db() -> ProjectDatabase { db } +/// The rendered MIR of the user function `name`, out of [`render_mir`]'s output. +fn function_mir<'a>(output: &'a str, name: &str) -> &'a str { + output + .split(&format!("fn user.{name}(")) + .nth(1) + .and_then(|tail| tail.split("\nfn ").next()) + .unwrap_or_else(|| panic!("no MIR for `user.{name}`:\n{output}")) +} + +/// Whether a line of MIR is a structural type test against a list or map type: +/// an `is_type` (which renders the type, `int[]`) or a `narrow_bind` (which +/// renders the type test, `List(Int)`). +fn is_structural_container_test(line: &str) -> bool { + (line.contains("is_type(") || line.contains("narrow_bind")) + && ["[]", "map<", "List(", "Map {"] + .iter() + .any(|spelling| line.contains(spelling)) +} + /// Lower all functions in a file to MIR and pretty-print them. fn render_mir(db: &ProjectDatabase, file: baml_base::SourceFile) -> String { // Dump in source order (by declaration span) — intrinsic and @@ -425,11 +444,7 @@ fn generic_class_destructure_field_projection_uses_instantiated_type() { // Scope the check to `user.f`'s body — auto-derived `to_json` / // `from_json` methods on `Box` legitimately have `void` locals // because T isn't instantiated in the auto-derive body. - let f_body = output - .split("fn user.f(") - .nth(1) - .and_then(|tail| tail.split("\nfn ").next()) - .unwrap_or(&output); + let f_body = function_mir(&output, "f"); assert!( !f_body .lines() @@ -438,6 +453,66 @@ fn generic_class_destructure_field_projection_uses_instantiated_type() { ); } +/// `json[]` and `map` are a `json` scrutinee's only list and map +/// members, so arms of exactly those types test only the runtime tag, also +/// when they bind the value (a narrowing arm tests a snapshot of the +/// scrutinee). +#[test] +fn json_container_arms_test_only_the_runtime_tag() { + let mut db = make_db(); + let file = db.file( + "test.baml", + r#" + function size(value: json) -> int { + match (value) { + let object: map => object.length(), + let items: json[] => items.length(), + _ => 0, + } + } + "#, + ); + let output = render_mir(&db, file); + let body = function_mir(&output, "size"); + assert!( + body.contains(", MAP);") && body.contains(", LIST);"), + "the json container arms do not test the runtime tag:\n{body}" + ); + assert!( + !body.lines().any(is_structural_container_test), + "a json container arm kept its structural test:\n{body}" + ); +} + +/// An element-specific arm stays element-precise on a `json` scrutinee: its +/// list member is `json[]`, so the `LIST` tag alone cannot prove `int[]`. +#[test] +fn json_scrutinee_keeps_element_specific_arms_structural() { + let mut db = make_db(); + let file = db.file( + "test.baml", + r#" + function ints(value: json) -> int { + match (value) { + let xs: int[] => xs.length(), + _ => 0, + } + } + "#, + ); + let output = render_mir(&db, file); + let body = function_mir(&output, "ints"); + assert!( + body.lines().any(|line| is_structural_container_test(line) + && (line.contains("int[]") || line.contains("List(Int)"))), + "the `int[]` arm lost its structural test:\n{body}" + ); + assert!( + !body.contains(", LIST);"), + "the `int[]` arm tests only the LIST tag:\n{body}" + ); +} + #[test] fn match_or_mixed_array_class_binding_uses_branch_local_rest_type() { let mut db = make_db(); diff --git a/baml_language/crates/baml_tests/tests/llm_call_memory.rs b/baml_language/crates/baml_tests/tests/llm_call_memory.rs index 35fe12c4b49..54dd491693e 100644 --- a/baml_language/crates/baml_tests/tests/llm_call_memory.rs +++ b/baml_language/crates/baml_tests/tests/llm_call_memory.rs @@ -3,9 +3,14 @@ //! A Rust test because BAML cannot observe allocations. Each parse used to build //! a schema-aligned-parsing model of every type in the program, the standard //! library included: about 5 MB per call for a one-function program, kept alive -//! until the VM collected it. A counting global allocator measures every byte -//! allocated while the calls run, on every thread, so the bound does not depend -//! on when the GC runs. +//! until the VM collected it. And every `match` type test that failed +//! canonicalized the recursive `json` alias: the request builder tests each +//! node of the request JSON that way, so the cost grew with the conversation. +//! Parsing a reply also used to build a second, external copy of the parsed +//! value and land it on the heap by name, so the parse cost of a large or +//! streamed output was larger than it needs to be. +//! A counting global allocator measures every byte allocated while the calls +//! run, on every thread, so the bounds do not depend on when the GC runs. #![allow( unsafe_code, @@ -74,12 +79,28 @@ static GLOBAL: CountingAllocator = CountingAllocator; const WARM_UP_CALLS: usize = 3; /// Calls measured. const MEASURED_CALLS: usize = 20; -/// The ceiling for one small LLM call. A debug build allocated about 3.4 MB per -/// call with the fix and about 13 MB before it. -const MAX_BYTES_PER_CALL: usize = 8 * 1024 * 1024; +/// The ceiling for one small LLM call. A debug build allocates about 0.7 MB; it +/// allocated about 3.4 MB while failed type tests expanded `json`, and about +/// 13 MB while each parse converted the whole program. +const MAX_BYTES_PER_CALL: usize = 2 * 1024 * 1024; /// The extra bytes per call that 400 unrelated classes may cost: about 22 KB with /// the fix, about 6.3 MB before it (each call converted every class twice). const MAX_UNRELATED_GROWTH_PER_CALL: usize = 256 * 1024; +/// User/assistant exchanges in the long conversation. +const LONG_CONVERSATION_TURNS: usize = 10; +/// Items in the large parsed output, and the characters a stream adds per batch. +const PARSED_ITEMS: i64 = 200; +const STREAM_STEP: usize = 256; +/// The ceiling for parsing a 13 KB reply of 200 items: about 665 KB; about +/// 1.15 MB while each parse built an external copy of the value. +const MAX_FINAL_PARSE_BYTES: usize = 900 * 1024; +/// The ceiling for the same reply streamed in 52 batches, each reparsed from +/// the start: about 32.3 MB; about 45.4 MB with the external copy. +const MAX_STREAMED_PARSE_BYTES: usize = 40 * 1024 * 1024; +/// The extra bytes per call that one more prompt message may cost: about 35 KB, +/// and about 1.56 MB while every failed type test on the request JSON expanded +/// `json`. +const MAX_BYTES_PER_MESSAGE: usize = 128 * 1024; fn completed_json(text: &str) -> String { let text = serde_json::to_string(text).expect("response text is JSON-serializable"); @@ -110,14 +131,19 @@ client Fake = openai.ResponsesClient.new(model = "fake", api_key = "test-key", b "#, base_url = server.uri() ); + (server, engine(&source)) +} + +/// An engine running `source` with the native sys ops. +fn engine(source: &str) -> Arc { let engine = BexEngine::new_with_runtime_compiler( - compile_source(&source), + compile_source(source), Arc::new(sys_ops::SysOps::native()), Vec::new(), bex_project::runtime_compiler(), ) .expect("engine builds"); - (server, Arc::new(engine)) + Arc::new(engine) } async fn call(engine: &Arc, function: &str) -> BexExternalValue { @@ -163,6 +189,72 @@ fn unrelated_classes(count: usize) -> String { .collect() } +/// A chat program whose prompt holds `turns` user/assistant exchanges before +/// the last question. +fn chat_program(turns: usize) -> String { + let history: String = (0..turns) + .map(|turn| { + format!("${{role(\"user\")}}Question {turn}?${{role(\"assistant\")}}Answer {turn}.") + }) + .collect(); + format!( + r#" +function Chat() -> string {{ + client: Fake + prompt: `${{role("system")}}Be concise.{history}${{role("user")}}Say hi.` +}} + +function main() -> string {{ + Chat() +}} +"# + ) +} + +/// A program that parses a 13 KB JSON reply of `PARSED_ITEMS` items: whole in +/// `final_parse`, and in `stream_parse` the way `ai.stream.Stream` does, one +/// growing prefix per batch through one parse cache. +fn parse_program() -> String { + let items: Vec = (0..PARSED_ITEMS) + .map(|i| format!(r#"{{"id": {i}, "name": "item number {i}", "tags": ["alpha", "beta"]}}"#)) + .collect(); + let text = serde_json::to_string(&format!("[{}]", items.join(", "))) + .expect("the reply is JSON-serializable"); + format!( + r#" +class Item {{ + id: int, + name: string, + tags: string[], +}} + +function final_parse() -> int {{ + baml.sap.parse({text}).length() +}} + +function stream_parse() -> int {{ + let text = {text}; + let cache = baml.sap._new_parse_cache(); + let yielded = 0; + let end = {STREAM_STEP}; + while (end < text.length() + {STREAM_STEP}) {{ + let upto = if (end > text.length()) {{ text.length() }} else {{ end }}; + let parsed: Item[] | baml.sap._NoYield = cache._parse_partial(text.slice(0, upto)); + match (parsed) {{ + let items: Item[] => {{ + // Only a prefix shorter than the reply shows a partial value. + if (upto < text.length()) {{ yielded = yielded + 1; }} + }}, + _ => {{}}, + }}; + end = end + {STREAM_STEP}; + }} + yielded +}} +"# + ) +} + const PICK_TOOL: &str = r#" class Choice { tool string @@ -224,3 +316,76 @@ function main() -> string { "one LLM call allocated {bytes} bytes (limit {MAX_BYTES_PER_CALL})" ); } + +#[tokio::test] +async fn call_cost_does_not_grow_with_conversation_length() { + let _measuring = MEASURING.lock().await; + let (_short_server, short) = engine_for(&chat_program(0), "hi").await; + let (_long_server, long) = engine_for(&chat_program(LONG_CONVERSATION_TURNS), "hi").await; + let short_bytes = bytes_per_call(&short, "main").await; + let long_bytes = bytes_per_call(&long, "main").await; + let extra_messages = 2 * LONG_CONVERSATION_TURNS; + let per_message = long_bytes.saturating_sub(short_bytes) / extra_messages; + assert!( + per_message <= MAX_BYTES_PER_MESSAGE, + "each prompt message added {per_message} bytes per call ({short_bytes} -> {long_bytes} \ + for {extra_messages} more messages); building the request must stay linear in its size \ + with a small constant" + ); +} + +#[tokio::test] +async fn large_parse_allocates_little() { + let _measuring = MEASURING.lock().await; + let engine = engine(&parse_program()); + assert_eq!( + call(&engine, "final_parse").await, + BexExternalValue::Int(PARSED_ITEMS) + ); + let bytes = bytes_per_call(&engine, "final_parse").await; + assert!( + bytes <= MAX_FINAL_PARSE_BYTES, + "parsing {PARSED_ITEMS} items allocated {bytes} bytes (limit {MAX_FINAL_PARSE_BYTES})" + ); +} + +#[tokio::test] +async fn streamed_parse_allocates_little() { + let _measuring = MEASURING.lock().await; + let engine = engine(&parse_program()); + assert!( + matches!(call(&engine, "stream_parse").await, BexExternalValue::Int(yields) if yields > 0), + "the stream never yielded a partial value" + ); + let bytes = bytes_per_call(&engine, "stream_parse").await; + assert!( + bytes <= MAX_STREAMED_PARSE_BYTES, + "streaming {PARSED_ITEMS} items allocated {bytes} bytes (limit {MAX_STREAMED_PARSE_BYTES})" + ); +} + +/// The parser drops a `@skip` field, so the types only a skipped field names +/// cost a parse nothing. +#[tokio::test] +async fn call_cost_does_not_grow_with_types_only_skipped_fields_name() { + let _measuring = MEASURING.lock().await; + let (_small_server, small) = engine_for(PICK_TOOL, r#"{"tool": "search"}"#).await; + let skipping_program = format!( + "{}\n{}", + unrelated_classes(400), + PICK_TOOL.replace( + "class Choice {\n tool string\n}", + "class Choice {\n tool string\n extra Unrelated399 @skip\n}", + ) + ); + assert!(skipping_program.contains("extra Unrelated399 @skip")); + let (_skipping_server, skipping) = engine_for(&skipping_program, r#"{"tool": "search"}"#).await; + let small_bytes = bytes_per_call(&small, "main").await; + let skipping_bytes = bytes_per_call(&skipping, "main").await; + let growth = skipping_bytes.saturating_sub(small_bytes); + assert!( + growth <= MAX_UNRELATED_GROWTH_PER_CALL, + "400 classes a skipped field names added {growth} bytes per call \ + ({small_bytes} -> {skipping_bytes}); the parse must not convert them" + ); +} diff --git a/baml_language/crates/baml_type/src/normalize.rs b/baml_language/crates/baml_type/src/normalize.rs index 9a82584a24e..6b545217525 100644 --- a/baml_language/crates/baml_type/src/normalize.rs +++ b/baml_language/crates/baml_type/src/normalize.rs @@ -414,17 +414,18 @@ pub trait TypeContext { if sub == sup { return true; } - // Narrow definite-mismatch filter. Unlike `equivalent`, differing heads - // do NOT generally refute subtyping (a literal is a subtype of its base, - // a member of its union, a class of an interface it implements) — but - // those pairs are different `Ty` *variants*, for which - // `heads_definitely_differ` already answers `false`. Of the same-variant - // nominal pairs it does claim, only interface-to-interface can still be - // a subtype under different names (`A <: B` iff `A` requires `B`), so it - // is excluded. For the rest the canonical walk below is a foregone - // `false`: class heads are preserved by canonicalization and the class - // subtype rule requires equal names (generic args are invariant), while - // enums and enum variants are nominal with no cross-name rule at all. + // Definite-mismatch filter. Unlike `equivalent`, differing heads do NOT + // generally refute subtyping (a literal is a subtype of its base, a + // member of its union, a class of an interface it implements) — but + // `heads_definitely_differ` claims none of those pairs: a literal and + // its base share a category, and unions and interfaces have none. Of + // the same-variant nominal pairs it does claim, only + // interface-to-interface can still be a subtype under different names + // (`A <: B` iff `A` requires `B`), so it is excluded. For the rest the + // canonical walk below is a foregone `false`: class heads are preserved + // by canonicalization and the class subtype rule requires equal names + // (generic args are invariant), enums and enum variants are nominal + // with no cross-name rule at all, and no rule relates two categories. if !matches!((sub, sup), (Ty::Interface(..), Ty::Interface(..))) && heads_definitely_differ(sub, sup) { @@ -624,24 +625,34 @@ pub fn is_subtype>(sub: &Ty, sup: &Ty, ctx: &C) /// are rewritten), and equality is nominal, so two of the same kind with /// different names can never share a canonical form. /// -/// This is deliberately *conservative* — it decides only same-kind pairs whose -/// heads are stable and never uses a cross-kind "different discriminant ⇒ -/// differ" catch-all. That catch-all would be unsound under the current -/// normalization: `TypeAlias` / `Union` / `AssociatedTypeProjection` / generic `Media` / -/// `Infer` heads can be rewritten into any other shape. Leaving every such case -/// undecided (`false`) preserves correctness; only the unambiguous nominal -/// misses are fast-rejected. Context-independent: canonicalization preserves -/// these nominal heads regardless of the `TypeContext`. +/// Across kinds it decides only pairs whose heads both have a [`Category`] +/// (see [`Category::of_head`]), never with a "different discriminant ⇒ differ" +/// catch-all. That catch-all would be unsound: `TypeAlias` / `Union` / +/// `AssociatedTypeProjection` heads can be rewritten into any other shape, so +/// they have no category and stay undecided (`false`). Canonicalization keeps +/// every head that has a category, so two different categories never share a +/// canonical form — and no subtyping rule relates them either: no concrete type +/// is a subtype of another (TYPE_SYSTEM.md §Concrete Types), and a literal and +/// its base, or a variant and its enum, share a category. This is the common +/// miss of a runtime `match` on a `json` value: a string or a list tested +/// against `map` would otherwise expand the recursive alias on +/// every test. Context-independent: canonicalization preserves these heads +/// regardless of the `TypeContext`. fn heads_definitely_differ(a: &Ty, b: &Ty) -> bool { match (a, b) { (Ty::Class(q1, ..), Ty::Class(q2, ..)) | (Ty::Interface(q1, ..), Ty::Interface(q2, ..)) | (Ty::Enum(q1, ..), Ty::Enum(q2, ..)) => q1 != q2, (Ty::EnumVariant(q1, v1, ..), Ty::EnumVariant(q2, v2, ..)) => q1 != q2 || v1 != v2, - _ => false, + _ => categories_differ(Category::of_head(a), Category::of_head(b)), } } +/// Whether two heads are of different, known categories. +fn categories_differ(a: Option, b: Option) -> bool { + matches!((a, b), (Some(a), Some(b)) if a != b) +} + /// Free-function form of [`TypeContext::definitely_disjoint`], for a context held /// by value. Pending removal once every caller uses the method form. pub fn definitely_disjoint>(a: &Ty, b: &Ty, ctx: &C) -> bool { @@ -773,10 +784,11 @@ impl NormalTy { // CONCRETE-TYPE DISJOINTNESS // ═══════════════════════════════════════════════════════════════════════════ -/// Top-level concrete category of a ground-headed type, used only for the -/// cross-category check (a value of one category is never a value of another). -/// Same-category pairs are decided by the structural arms of -/// [`NormalTy::is_disjoint_from`], not by this. +/// Top-level concrete category of a ground-headed type, for the cross-category +/// checks: a value of one category is never a value of another +/// ([`NormalTy::is_disjoint_from`]), and no subtyping rule or canonical form +/// relates two categories ([`heads_definitely_differ`]). Same-category pairs +/// are decided by the structural arms, not by this. #[derive(PartialEq, Eq)] enum Category { Int, @@ -1017,6 +1029,47 @@ impl NormalTy { } } +/// The one table behind [`Category::of_head`] and +/// [`Category::of_interned_head`]. Both are `ty_family!` members with the same +/// variant names, so `$member` is the enum's path and `$none` adds the member's own +/// variants that have no category. Exhaustive on purpose: a new variant must +/// be classified here. +macro_rules! family_head_category { + ($ty:expr, $($member:ident)::+ $(, $none:pat)*) => { + Some(match $ty { + $($member)::+::Int => Category::Int, + $($member)::+::Bigint => Category::Bigint, + $($member)::+::Float => Category::Float, + $($member)::+::String => Category::String, + $($member)::+::Bool => Category::Bool, + $($member)::+::Literal(lit, _) => Category::of_literal(lit), + $($member)::+::Null => Category::Null, + $($member)::+::Uint8Array => Category::Uint8Array, + $($member)::+::Media(kind) => Category::Media(*kind), + $($member)::+::Void => Category::Void, + $($member)::+::RustType => Category::RustType, + $($member)::+::Type => Category::Type, + $($member)::+::Resource => Category::Resource, + $($member)::+::PromptAst => Category::PromptAst, + $($member)::+::Class(..) => Category::Class, + $($member)::+::List(_) => Category::List, + $($member)::+::Map { .. } => Category::Map, + $($member)::+::Enum(_) | $($member)::+::EnumVariant(..) => Category::Enum, + $($member)::+::Function { .. } => Category::Function, + $($member)::+::Future(..) => Category::Future, + $($member)::+::Interface(..) + | $($member)::+::Union(_) + | $($member)::+::TypeAlias(_) + | $($member)::+::TypeVar(_) + | $($member)::+::AssociatedTypeProjection { .. } + | $($member)::+::Unknown + | $($member)::+::Never + | $($member)::+::Error + $(| $none)* => return None, + }) + }; +} + impl Category { fn of_literal(lit: &Literal) -> Category { match lit { @@ -1027,6 +1080,22 @@ impl Category { Literal::Bool(_) => Category::Bool, } } + + /// The category of a head that canonicalization keeps, read before + /// canonicalizing: the [`Ty`] twin of [`NormalTy::head_category`]. `None` + /// for a head canonicalization can rewrite into another kind (an alias, a + /// union, a projection), one a context-driven rule relates across kinds + /// (an interface, a type variable), and the sentinels (`unknown`, `never`, + /// the error type). + fn of_head(ty: &Ty) -> Option { + family_head_category!(ty, Ty) + } + + /// [`Category::of_head`] for the interned kind. A live inference variable + /// can still become any type, so it has no category. + fn of_interned_head(ty: &interned::InferTy) -> Option { + family_head_category!(ty, interned::InferTy, interned::InferTy::InferVar { .. }) + } } // ═══════════════════════════════════════════════════════════════════════════ @@ -3166,7 +3235,7 @@ fn interned_heads_definitely_differ(a: &interned::Ty, b: &interned::Ty) -> bool | (K::Interface(q1, ..), K::Interface(q2, ..)) | (K::Enum(q1, ..), K::Enum(q2, ..)) => q1 != q2, (K::EnumVariant(q1, v1, ..), K::EnumVariant(q2, v2, ..)) => q1 != q2 || v1 != v2, - _ => false, + (a, b) => categories_differ(Category::of_interned_head(a), Category::of_interned_head(b)), } } diff --git a/baml_language/crates/baml_type/src/normalize/tests.rs b/baml_language/crates/baml_type/src/normalize/tests.rs index 09a3f303a75..b69f65ff0c0 100644 --- a/baml_language/crates/baml_type/src/normalize/tests.rs +++ b/baml_language/crates/baml_type/src/normalize/tests.rs @@ -1,7 +1,7 @@ //! Unit tests for the canonical type algebra, driven by a hand-built //! [`TypeContext`] stub that records nominal facts directly. -use std::collections::HashMap; +use std::{cell::Cell, collections::HashMap}; use super::*; use crate::{DeclName, Freshness, FunctionParamTy, Literal, Name, Ty}; @@ -25,6 +25,9 @@ struct Ctx { /// `(base, interface head, member) → reduced type` projection facts, for the /// `project` oracle. A `Vec` (not a map) because `Ty` is not `Hash`. projections: Vec<(Ty, DeclName, Name, Ty)>, + /// How many times the algebra asked for an alias definition: a refutation + /// that skips canonicalization never asks. + alias_lookups: Cell, } fn nominal_head(ty: &Ty) -> Option { @@ -55,6 +58,7 @@ impl TypeContext for Ctx { } fn alias_def(&self, name: &DeclName) -> Option { + self.alias_lookups.set(self.alias_lookups.get() + 1); self.aliases.get(name).cloned() } @@ -1457,6 +1461,166 @@ fn any_function_existentials_are_covariant_in_their_pins() { assert!(!is_subtype(&other, &any_function(vec![]), &ctx)); } +// ── head-category refutation ─────────────────────────────────────────────-- + +/// `type Json = null | bool | int | float | string | Json[] | map` +/// registered under the given name: the recursive alias whose canonical form +/// is expensive to build. +fn json_alias(ctx: &mut Ctx, name: &str) { + ctx.aliases.insert( + qtn(name), + union(vec![ + Ty::null(), + Ty::bool(), + Ty::int(), + Ty::float(), + Ty::string(), + list(alias(name)), + map_ty(Ty::string(), alias(name)), + ]), + ); +} + +/// One type of every `Ty` variant, plus the facts that make the context-driven +/// rules fire: `Box` and `int` implement `Displayable`, `T: Displayable`, +/// `(Box as Displayable).Item` reduces to `int`, and `Json` is recursive. +fn kind_corpus() -> (Ctx, Vec) { + let mut ctx = Ctx::default(); + json_alias(&mut ctx, "Json"); + ctx.enums + .insert(qtn("Side"), vec![Name::new("L"), Name::new("R")]); + ctx.impls.push((qtn("Box"), qtn("Displayable"))); + ctx.prim_impls.push(("int", qtn("Displayable"))); + ctx.var_bounds + .insert(param(0, "T"), vec![iface("Displayable")]); + ctx.projections.push(( + class("Box"), + qtn("Displayable"), + Name::new("Item"), + Ty::int(), + )); + let corpus = vec![ + Ty::int(), + bigint(), + Ty::float(), + Ty::string(), + Ty::bool(), + Ty::null(), + Ty::Uint8Array, + Ty::Media(MediaKind::Image), + Ty::Media(MediaKind::Generic), + Ty::Void, + Ty::RustType, + Ty::Type, + Ty::Resource, + Ty::PromptAst, + lit_int(1), + lit_float("1.5"), + lit_str("user"), + lit_bool(true), + class("Box"), + class1("Gen", Ty::int()), + enum_ty("Side"), + variant("Side", "L"), + list(Ty::int()), + list(alias("Json")), + map_ty(Ty::string(), alias("Json")), + map_ty(Ty::string(), Ty::int()), + Ty::Function { + params: Box::new([FunctionParamTy::required(None, Ty::int())]), + ret: Box::new(Ty::bool()), + throws: Box::new(Ty::Never), + }, + Ty::Future(Box::new(Ty::int()), Box::new(Ty::Never)), + union(vec![Ty::int(), Ty::string()]), + Ty::optional(Ty::string()), + alias("Json"), + iface("Displayable"), + typevar(0, "T"), + projection(class("Box"), "Displayable", "Item"), + Ty::Unknown, + Ty::Never, + Ty::Error, + ]; + (ctx, corpus) +} + +/// The head-category refutation is only a shortcut: over every pair of +/// variants, `is_subtype` and `equivalent` give the verdict the canonical +/// relation gives. +#[test] +fn subtype_and_equivalence_agree_with_the_canonical_relation_for_every_kind_pair() { + let (ctx, corpus) = kind_corpus(); + for sub in &corpus { + for sup in &corpus { + let canonical_sub = NormalTy::canonical(sub, &ctx); + let canonical_sup = NormalTy::canonical(sup, &ctx); + assert_eq!( + ctx.is_subtype(sub, sup), + canonical_sub.is_subtype_of(&canonical_sup, &ctx, &mut HashSet::new()), + "subtype verdict diverged for {sub:?} <: {sup:?}" + ); + assert_eq!( + ctx.equivalent(sub, sup), + canonical_sub == canonical_sup, + "equivalence verdict diverged for {sub:?} == {sup:?}" + ); + } + } +} + +/// Pairs whose top-level kinds no rule relates, the ones a runtime `match` on +/// a `json` value asks about for every node: none of them may canonicalize +/// (and so expand) the recursive alias. +fn cross_category_pairs() -> Vec<(Ty, Ty)> { + let json_map = map_ty(Ty::string(), alias("Json")); + let json_list = list(alias("Json")); + vec![ + (lit_str("user"), json_map.clone()), + (lit_str("user"), json_list.clone()), + (json_list.clone(), json_map.clone()), + (json_map.clone(), json_list.clone()), + (Ty::int(), json_map.clone()), + (lit_int(1), Ty::float()), + (Ty::null(), json_list), + (class("Box"), json_map.clone()), + (enum_ty("Side"), json_map), + ] +} + +#[test] +fn cross_category_pairs_are_refuted_without_canonicalizing() { + let mut ctx = Ctx::default(); + json_alias(&mut ctx, "Json"); + for (sub, sup) in cross_category_pairs() { + assert!(!ctx.is_subtype(&sub, &sup), "{sub:?} <: {sup:?}"); + assert!(!ctx.equivalent(&sub, &sup), "{sub:?} == {sup:?}"); + assert_eq!( + ctx.alias_lookups.get(), + 0, + "{sub:?} vs {sup:?} canonicalized an alias" + ); + } +} + +/// Kinds a canonical form can rewrite, or that a context-driven rule relates +/// across kinds, are never refuted by their head alone. +#[test] +fn rewritable_and_context_driven_kinds_are_not_refuted() { + let (ctx, _) = kind_corpus(); + let json_map = map_ty(Ty::string(), alias("Json")); + assert!(ctx.is_subtype(&lit_str("user"), &alias("Json"))); + assert!(ctx.is_subtype(&json_map, &alias("Json"))); + assert!(ctx.is_subtype(&Ty::int(), &iface("Displayable"))); + assert!(ctx.is_subtype(&class("Box"), &iface("Displayable"))); + assert!(ctx.is_subtype(&typevar(0, "T"), &iface("Displayable"))); + assert!(ctx.is_subtype(&projection(class("Box"), "Displayable", "Item"), &Ty::int())); + assert!(ctx.is_subtype(&Ty::Never, &json_map)); + assert!(ctx.is_subtype(&json_map, &Ty::Unknown)); + assert!(ctx.is_subtype(&Ty::Error, &json_map)); + assert!(ctx.is_subtype(&lit_str("user"), &Ty::optional(Ty::string()))); +} + // ── interned entry (S4b) ─────────────────────────────────────────────────── mod interned_entry { @@ -1524,6 +1688,62 @@ mod interned_entry { )); } + /// The interned entries give the canonical verdict for every kind pair too. + #[test] + fn interned_verdicts_agree_with_the_canonical_relation_for_every_kind_pair() { + let (ctx, corpus) = kind_corpus(); + let corpus: Vec<_> = corpus.iter().map(it).collect(); + let cache = InternedCanonicalCache::default(); + for sub in &corpus { + for sup in &corpus { + let canonical_sub = NormalTy::canonical_interned(sub, &ctx); + let canonical_sup = NormalTy::canonical_interned(sup, &ctx); + let subtype = + canonical_sub.is_subtype_of(&canonical_sup, &ctx, &mut HashSet::new()); + let equivalent = canonical_sub == canonical_sup; + assert_eq!( + is_subtype_interned(sub, sup, &ctx), + subtype, + "{sub:?} <: {sup:?}" + ); + assert_eq!( + cache.is_subtype(sub, sup, &ctx), + subtype, + "{sub:?} <: {sup:?}" + ); + assert_eq!( + equivalent_interned(sub, sup, &ctx), + equivalent, + "{sub:?} == {sup:?}" + ); + assert_eq!( + cache.equivalent(sub, sup, &ctx), + equivalent, + "{sub:?} == {sup:?}" + ); + } + } + } + + #[test] + fn interned_cross_category_pairs_are_refuted_without_canonicalizing() { + let mut ctx = Ctx::default(); + json_alias(&mut ctx, "Json"); + let cache = InternedCanonicalCache::default(); + for (sub, sup) in cross_category_pairs() { + let (sub, sup) = (it(&sub), it(&sup)); + assert!(!is_subtype_interned(&sub, &sup, &ctx), "{sub:?} <: {sup:?}"); + assert!(!equivalent_interned(&sub, &sup, &ctx), "{sub:?} == {sup:?}"); + assert!(!cache.is_subtype(&sub, &sup, &ctx), "{sub:?} <: {sup:?}"); + assert!(!cache.equivalent(&sub, &sup, &ctx), "{sub:?} == {sup:?}"); + assert_eq!( + ctx.alias_lookups.get(), + 0, + "{sub:?} vs {sup:?} canonicalized an alias" + ); + } + } + #[test] fn canonical_cache_matches_reject_free_canonical_relations() { let mut ctx = Ctx::default(); diff --git a/baml_language/crates/bex_engine/src/lib.rs b/baml_language/crates/bex_engine/src/lib.rs index 646b9de35a9..b7218484b42 100644 --- a/baml_language/crates/bex_engine/src/lib.rs +++ b/baml_language/crates/bex_engine/src/lib.rs @@ -2248,43 +2248,7 @@ impl BexEngine { if let Object::Class(cls) = obj { defs.insert( ::sys_types::DefKey::new(cls.type_tag, cls.name.clone()), - sys_types::ClassDefinition { - name: cls.name.display_name().to_string(), - description: cls.description.clone(), - docstring: cls.docstring.clone(), - alias: cls.alias.clone(), - stream_done: cls.stream_done, - fields: cls - .fields - .iter() - .map(|f| sys_types::ClassFieldDefinition { - name: f.name.clone(), - field_type: f - .field_type - .try_map_heads(&mut bex_vm_types::TypeHead::to_tagged_name) - .unwrap_or_else(|_| { - unreachable!( - "compiled class fields name compiled declarations" - ) - }), - field_template: Some( - f.field_template - .try_map_heads(&mut bex_vm_types::TypeHead::to_tagged_name) - .unwrap_or_else(|_| { - unreachable!( - "compiled class fields name compiled declarations" - ) - }), - ), - description: f.description.clone(), - docstring: f.docstring.clone(), - alias: f.alias.clone(), - skip: f.skip, - stream_done: f.stream_done, - must_exist: f.must_exist, - }) - .collect(), - }, + bex_vm::definitions::class_definition(cls), ); } } @@ -2302,109 +2266,18 @@ impl BexEngine { if let Object::Enum(enm) = obj { defs.insert( ::sys_types::DefKey::new(enm.type_tag, enm.name.clone()), - sys_types::EnumDefinition { - name: enm.name.display_name().to_string(), - description: enm.description.clone(), - docstring: enm.docstring.clone(), - alias: enm.alias.clone(), - variants: enm - .variants - .iter() - .filter(|v| !v.skip) - .map(|v| sys_types::EnumVariantDefinition { - name: v.name.clone(), - description: v.description.clone(), - docstring: v.docstring.clone(), - alias: v.alias.clone(), - }) - .collect(), - }, + bex_vm::definitions::enum_definition(enm), ); } } defs } - /// Carry a declaration's field type onto the sys-op lane. - /// - /// Total by invariant: a live declaration's heads are resolved pointers to - /// declarations, so each yields its identity and its own name. - fn lane_ty(ty: &bex_vm_types::RuntimeTy) -> ::sys_types::SapTy { - ty.try_map_heads(&mut bex_vm_types::TypeHead::to_tagged_name) - .unwrap_or_else(|head| { - unreachable!("a live declaration's field names a declaration: {head}") - }) - } - - /// [`Self::lane_ty`] for a field's symbolic template, same totality - /// argument: every head a live declaration's template names is resolved. - fn lane_template(template: &bex_vm_types::TyTemplate) -> ::sys_types::SapTyTemplate { - template - .try_map_heads(&mut bex_vm_types::TypeHead::to_tagged_name) - .unwrap_or_else(|head| { - unreachable!("a live declaration's field template names a declaration: {head}") - }) - } - - fn enum_definition(enm: &bex_vm_types::Enum) -> sys_types::EnumDefinition { - sys_types::EnumDefinition { - name: enm.name.display_name().to_string(), - description: enm.description.clone(), - docstring: enm.docstring.clone(), - alias: enm.alias.clone(), - variants: enm - .variants - .iter() - .filter(|variant| !variant.skip) - .map(|variant| sys_types::EnumVariantDefinition { - name: variant.name.clone(), - description: variant.description.clone(), - docstring: variant.docstring.clone(), - alias: variant.alias.clone(), - }) - .collect(), - } - } - - fn class_definition( - class: &bex_vm_types::Class, - _permit: bex_heap::PermitProof<'_>, - ) -> sys_types::ClassDefinition { - sys_types::ClassDefinition { - name: class.name.display_name().to_string(), - description: class.description.clone(), - docstring: class.docstring.clone(), - alias: class.alias.clone(), - stream_done: class.stream_done, - fields: class - .fields - .iter() - .filter(|field| !field.skip) - .map(|field| sys_types::ClassFieldDefinition { - name: field.name.clone(), - field_type: Self::lane_ty(&field.field_type), - field_template: Some(Self::lane_template(&field.field_template)), - description: field.description.clone(), - docstring: field.docstring.clone(), - alias: field.alias.clone(), - skip: field.skip, - stream_done: field.stream_done, - must_exist: field.must_exist, - }) - .collect(), - } - } - /// Gather runtime definitions from the type descriptors passed directly /// to a sys-op. Type arguments are lowered as ordinary `Object::Type` /// arguments, so this is the last synchronous chokepoint before the permit /// is released for async work. - fn runtime_type_overlay( - &self, - vm: &BexVm, - args: &[Value], - permit: bex_heap::PermitProof<'_>, - ) -> RuntimeTypeOverlay { + fn runtime_type_overlay(&self, vm: &BexVm, args: &[Value]) -> RuntimeTypeOverlay { let mut overlay = RuntimeTypeOverlay::default(); for value in args { let Some(type_ptr) = value.as_object_ptr() else { @@ -2425,7 +2298,7 @@ impl BexEngine { overlay .class_definitions .entry(head.clone()) - .or_insert_with(|| Self::class_definition(class, permit)); + .or_insert_with(|| bex_vm::definitions::class_definition(class)); overlay .class_handles .entry(head) @@ -2439,7 +2312,7 @@ impl BexEngine { overlay .enum_definitions .entry(head.clone()) - .or_insert_with(|| Self::enum_definition(enm)); + .or_insert_with(|| bex_vm::definitions::enum_definition(enm)); overlay .enum_handles .entry(head) @@ -2499,28 +2372,7 @@ impl BexEngine { }; classes.insert( ::sys_types::DefKey::new(class.type_tag, class.name.clone()), - sys_types::ClassDefinition { - name: class.name.display_name().to_string(), - description: class.description.clone(), - docstring: class.docstring.clone(), - alias: class.alias.clone(), - stream_done: class.stream_done, - fields: class - .fields - .iter() - .map(|field| sys_types::ClassFieldDefinition { - name: field.name.clone(), - field_type: Self::lane_ty(&field.field_type), - field_template: Some(Self::lane_template(&field.field_template)), - description: field.description.clone(), - docstring: field.docstring.clone(), - alias: field.alias.clone(), - skip: field.skip, - stream_done: field.stream_done, - must_exist: field.must_exist, - }) - .collect(), - }, + bex_vm::definitions::class_definition(class), ); named_owners.insert(class.name.to_string(), handle.clone()); } @@ -2530,23 +2382,7 @@ impl BexEngine { }; enums.insert( ::sys_types::DefKey::new(enm.type_tag, enm.name.clone()), - sys_types::EnumDefinition { - name: enm.name.display_name().to_string(), - description: enm.description.clone(), - docstring: enm.docstring.clone(), - alias: enm.alias.clone(), - variants: enm - .variants - .iter() - .filter(|variant| !variant.skip) - .map(|variant| sys_types::EnumVariantDefinition { - name: variant.name.clone(), - description: variant.description.clone(), - docstring: variant.docstring.clone(), - alias: variant.alias.clone(), - }) - .collect(), - }, + bex_vm::definitions::enum_definition(enm), ); named_owners.insert(enm.name.to_string(), handle.clone()); } @@ -5853,8 +5689,7 @@ impl BexEngine { // runs its sys-ops to completion. let op_cancel = thread.sysop_cancel.clone(); - let runtime_type_overlay = - self.runtime_type_overlay(&thread.vm, &args, thread.proof()); + let runtime_type_overlay = self.runtime_type_overlay(&thread.vm, &args); let runtime_compile_request = match operation { SysOp::ReflectPackageCompile => { Some(Ok(Self::runtime_compile_request(&thread.vm, &args)?)) diff --git a/baml_language/crates/bex_sap/Cargo.toml b/baml_language/crates/bex_sap/Cargo.toml index d9085f053eb..f392936d57b 100644 --- a/baml_language/crates/bex_sap/Cargo.toml +++ b/baml_language/crates/bex_sap/Cargo.toml @@ -18,7 +18,6 @@ workspace = true [dependencies] baml_type = { workspace = true } -bex_external_types = { workspace = true } bex_str = { workspace = true } bex_vm_types = { workspace = true } sys_types = { workspace = true } diff --git a/baml_language/crates/bex_sap/src/lib.rs b/baml_language/crates/bex_sap/src/lib.rs index d3dc4f5cd05..2db7436a6f4 100644 --- a/baml_language/crates/bex_sap/src/lib.rs +++ b/baml_language/crates/bex_sap/src/lib.rs @@ -2,7 +2,10 @@ use ouroboros::self_referencing; use sys_types::DefKey; pub use crate::jsonish::parse; -use crate::sap_model::{Ty, TypeRefDb}; +use crate::{ + deserializer::coercer::{ParsingContext, ParsingError}, + sap_model::{Ty, TyResolvedRef, TypeRefDb}, +}; pub mod baml_value; pub mod deserializer; @@ -10,7 +13,7 @@ pub mod jsonish; pub mod sap_model; #[cfg(test)] mod tests; -pub mod to_external; +pub mod to_baml_ty; #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum StreamingMode { @@ -43,14 +46,6 @@ impl CompiledSapModel { )?; Ok(Self { inner }) } - pub fn from_sys_op_context( - ctx: &::sys_types::SysOpContext, - target: sys_types::SapTy, - ) -> Result { - let type_ctx = sap_model::TypeCtx::for_target(ctx, &target); - Self::from_type_ctx(type_ctx, target) - } - pub fn db(&self) -> &TypeRefDb<'_, DefKey> { self.inner.borrow_db() } @@ -58,6 +53,18 @@ impl CompiledSapModel { pub fn ty(&self) -> &Ty<'_, DefKey> { self.inner.borrow_ty() } + + /// The parse target, resolved in [`Self::db`]. + /// + /// # Errors + /// + /// The target names a type the database does not hold: a schema the + /// converter accepted but cannot resolve, which no input text can fix. + pub fn resolved_target(&self) -> Result, ParsingError> { + self.db() + .resolve(self.ty()) + .map_err(|name| ParsingContext::new(self.db()).error_type_resolution(name)) + } } #[self_referencing] diff --git a/baml_language/crates/bex_sap/src/sap_model/convert.rs b/baml_language/crates/bex_sap/src/sap_model/convert.rs index d9ef52c1aed..d66aca35502 100644 --- a/baml_language/crates/bex_sap/src/sap_model/convert.rs +++ b/baml_language/crates/bex_sap/src/sap_model/convert.rs @@ -47,8 +47,9 @@ pub enum ConvertError { const MAX_RECURSION_DEPTH: usize = 16; -/// Contains stuff from [`sys_types::SysOpContext`] that we need for converting to the sap model: -/// the definitions one parse target reaches ([`TypeCtx::for_target`]), not the whole program's. +/// The class, enum, and type-alias definitions one parse target reaches, which the converter +/// needs to build the sap model. The caller reads them off the target's heap declarations, so a +/// parse pays only for the types its target uses, not for the whole program's. /// /// ## Representation /// - Unions should be flattened: @@ -71,9 +72,9 @@ pub struct TypeCtx { } impl TypeCtx { /// Builds the context from definitions that are closed under reference: - /// every name a class field or an alias body refers to is present too - /// ([`TypeCtx::for_target`] collects such a set). A missing name makes the - /// types that refer to it unparsable. + /// every name a class field or an alias body refers to is present too (the + /// declarations a type reaches on the heap are such a set). A missing name + /// makes the types that refer to it unparsable. pub fn new<'d>( class_definitions: impl IntoIterator, enum_definitions: IndexMap, @@ -149,38 +150,6 @@ impl TypeCtx { } } - /// The context for parsing into `target`: only the classes, enums, and type - /// aliases that `target` reaches (`reachable_definitions`), in program - /// order. Converting every definition in the program on each parse cost - /// several MB per LLM call; the reachable set is all the converter looks up. - pub fn for_target( - ctx: &::sys_types::SysOpContext, - target: &SapTy, - ) -> Self { - let reachable = - reachable_definitions(target, &ctx.class_definitions, &ctx.type_alias_definitions); - let enum_definitions = ctx - .enum_definitions - .iter() - .filter(|(name, _)| reachable.contains(*name)) - .map(|(name, definition)| (name.clone(), definition.clone())) - .collect(); - let type_alias_definitions = ctx - .type_alias_definitions - .iter() - .filter(|(name, _)| reachable.contains(*name)) - .map(|(name, ty)| (name.clone(), ty.clone())) - .collect(); - - Self::new( - ctx.class_definitions - .iter() - .filter(|(name, _)| reachable.contains(*name)), - enum_definitions, - &type_alias_definitions, - ) - } - /// Normalize a runtime-materialized parse target before converting it to /// the SAP model. Generic substitution happens after [`TypeCtx::new`] has /// simplified the declared class and alias definitions, so it can create a @@ -498,42 +467,6 @@ impl TypeCtx { } } -/// The definitions `target` can reach: the classes, interfaces, enums, and type -/// aliases it names, then, transitively, the ones named by a reached class's -/// field types (except `@skip` fields, which the parser drops) or a reached -/// alias's body. -/// -/// The walk is [`RuntimeTy::visit_heads`](baml_type::RuntimeTy::visit_heads), -/// the head walk `bex_vm::reachable` uses for heap declarations, run over the -/// context's definition tables (which the engine extracts from those same -/// declarations). Heads the converter never resolves, such as type arguments -/// and function types, are kept too. That only adds definitions, so every name -/// the converter looks up is present. -fn reachable_definitions( - target: &SapTy, - class_definitions: &IndexMap, - type_alias_definitions: &IndexMap, -) -> HashSet { - let mut reachable = HashSet::new(); - let mut pending = Vec::new(); - target.visit_heads(&mut |name: &DefKey| pending.push(name.clone())); - while let Some(name) = pending.pop() { - if !reachable.insert(name.clone()) { - continue; - } - if let Some(class) = class_definitions.get(&name) { - for field in class.fields.iter().filter(|field| !field.skip) { - field - .field_type - .visit_heads(&mut |head: &DefKey| pending.push(head.clone())); - } - } else if let Some(alias) = type_alias_definitions.get(&name) { - alias.visit_heads(&mut |head: &DefKey| pending.push(head.clone())); - } - } - reachable -} - fn check_parseable( name: &DefKey, class_definitions: &IndexMap, @@ -633,195 +566,3 @@ fn is_sap_parseable(ty: &SapTy) -> Result, ()> { | SapTy::Type => Err(()), } } - -#[cfg(test)] -mod tests { - use std::sync::Arc; - - use baml_type::{DeclarationName, Name, typetag::TypeTag}; - use sys_types::{ClassFieldDefinition, EnumVariantDefinition, SysOpContext}; - - use super::*; - - fn key(name: &str) -> DefKey { - DefKey::new( - TypeTag::fresh_dynamic(), - DeclarationName::Anonymous(Name::new(name)), - ) - } - - fn class(key: &DefKey, fields: Vec<(&str, SapTy)>) -> (DefKey, ClassDefinition) { - let fields = fields - .into_iter() - .map(|(name, field_type)| ClassFieldDefinition { - name: name.to_string(), - field_type, - field_template: None, - description: None, - docstring: None, - alias: None, - skip: false, - stream_done: false, - must_exist: false, - }) - .collect(); - let definition = ClassDefinition { - name: key.to_string(), - description: None, - docstring: None, - alias: None, - fields, - stream_done: false, - }; - (key.clone(), definition) - } - - fn enumeration(key: &DefKey) -> (DefKey, EnumDefinition) { - let definition = EnumDefinition { - name: key.to_string(), - description: None, - docstring: None, - alias: None, - variants: vec![EnumVariantDefinition { - name: "Only".to_string(), - description: None, - docstring: None, - alias: None, - }], - }; - (key.clone(), definition) - } - - fn class_ty(key: &DefKey) -> SapTy { - SapTy::Class(key.clone(), Box::new([])) - } - - /// A program with a reachable cluster around `Root` and types nothing in it - /// refers to. `Root -> Child -> Root` is a cycle; `Child` also names an enum - /// and a recursive alias whose body names `Leaf`. - struct Program { - root: DefKey, - child: DefKey, - leaf: DefKey, - color: DefKey, - tree: DefKey, - unrelated: DefKey, - unrelated_enum: DefKey, - ctx: SysOpContext, - } - - fn program() -> Program { - let [root, child, leaf, color, tree, unrelated, unrelated_enum] = [ - "Root", - "Child", - "Leaf", - "Color", - "Tree", - "Unrelated", - "UnrelatedEnum", - ] - .map(key); - let tree_body = - SapTy::union([SapTy::list(SapTy::TypeAlias(tree.clone())), class_ty(&leaf)]); - let mut ctx = SysOpContext::empty(); - ctx.class_definitions = Arc::new(IndexMap::from([ - class(&unrelated, vec![("value", SapTy::string())]), - class(&root, vec![("child", class_ty(&child))]), - class( - &child, - vec![ - ("parent", SapTy::optional(class_ty(&root))), - ("color", SapTy::Enum(color.clone())), - ("tree", SapTy::TypeAlias(tree.clone())), - ], - ), - class(&leaf, vec![("value", SapTy::int())]), - ])); - ctx.enum_definitions = Arc::new(IndexMap::from([ - enumeration(&unrelated_enum), - enumeration(&color), - ])); - ctx.type_alias_definitions = Arc::new(IndexMap::from([(tree.clone(), tree_body)])); - Program { - root, - child, - leaf, - color, - tree, - unrelated, - unrelated_enum, - ctx, - } - } - - fn names<'d, V: 'd>( - definitions: impl IntoIterator, - ) -> Vec<&'d DefKey> { - definitions.into_iter().map(|(name, _)| name).collect() - } - - /// The parser drops a `@skip` field, so a type only a skipped field names - /// is not part of the model. - #[test] - fn skipped_fields_reach_nothing() { - let (root, hidden) = (key("Root"), key("Hidden")); - let (root_key, mut root_definition) = class( - &root, - vec![("name", SapTy::string()), ("secret", class_ty(&hidden))], - ); - root_definition.fields[1].skip = true; - let mut ctx = SysOpContext::empty(); - ctx.class_definitions = Arc::new(IndexMap::from([ - (root_key, root_definition), - class(&hidden, vec![("value", SapTy::int())]), - ])); - - let type_ctx = TypeCtx::for_target(&ctx, &class_ty(&root)); - - assert_eq!(names(&type_ctx.class_definitions), vec![&root]); - } - - #[test] - fn primitive_target_keeps_no_definitions() { - let program = program(); - let ctx = TypeCtx::for_target(&program.ctx, &SapTy::string()); - assert!(ctx.class_definitions.is_empty()); - assert!(ctx.enum_definitions.is_empty()); - assert!(ctx.type_alias_definitions.is_empty()); - } - - #[test] - fn class_target_keeps_what_its_fields_reach_in_program_order() { - let program = program(); - let ctx = TypeCtx::for_target(&program.ctx, &class_ty(&program.root)); - - assert_eq!( - names(&ctx.class_definitions), - vec![&program.root, &program.child, &program.leaf], - ); - assert_eq!(names(&ctx.enum_definitions), vec![&program.color]); - assert_eq!(names(&ctx.type_alias_definitions), vec![&program.tree]); - assert!(!ctx.class_definitions.contains_key(&program.unrelated)); - assert!(!ctx.enum_definitions.contains_key(&program.unrelated_enum)); - } - - #[test] - fn alias_target_keeps_its_body() { - let program = program(); - let ctx = TypeCtx::for_target(&program.ctx, &SapTy::TypeAlias(program.tree.clone())); - - assert_eq!(names(&ctx.class_definitions), vec![&program.leaf]); - assert!(ctx.enum_definitions.is_empty()); - assert_eq!(names(&ctx.type_alias_definitions), vec![&program.tree]); - } - - #[test] - fn reachable_definitions_convert_without_unknown_names() { - let program = program(); - let root = class_ty(&program.root); - let ctx = TypeCtx::for_target(&program.ctx, &root); - let db = ctx.build_db().expect("the reachable set converts"); - let target = ctx.convert_ty(&root).expect("the target converts"); - assert!(db.resolve(&target).is_ok()); - } -} diff --git a/baml_language/crates/bex_sap/src/to_baml_ty.rs b/baml_language/crates/bex_sap/src/to_baml_ty.rs new file mode 100644 index 00000000000..a084b9b0aae --- /dev/null +++ b/baml_language/crates/bex_sap/src/to_baml_ty.rs @@ -0,0 +1,140 @@ +//! Conversion from SAP model types back to `SapTy`. +//! +//! A parsed value's metadata carries SAP model types; the VM, which builds the +//! value on the heap, turns them into the element, key, and value types of the +//! lists and maps it allocates. + +use sys_types::{DefKey, SapTy}; + +use crate::sap_model::{ + ArrayTy, ClassTy, EnumTy, EnumVariantTy, MapTy, MediaTy, StreamStateTy, Ty, TyResolvedRef, + TypeRefDb, UnionTy, +}; + +// ============================================================================ +// Type conversion: SAP types → SapTy +// ============================================================================ + +/// Convert a SAP type back to a `SapTy`. +/// +/// # Known simplifications (not round-trip safe) +/// +/// - `Optional(T)` was converted to `Union([Null, T])` on the way in; it stays as `Union` here. +/// - Aliases already resolved into the target type have no nominal name to recover; named alias +/// references inside containers are preserved. +/// - `EnumVariant` becomes `Enum` (variant specificity lost at the type level). +pub trait ToBamlTy { + fn to_baml_ty(&self, db: &TypeRefDb<'_, DefKey>) -> SapTy; +} + +impl ToBamlTy for TyResolvedRef<'_, DefKey> { + fn to_baml_ty(&self, db: &TypeRefDb<'_, DefKey>) -> SapTy { + match self { + TyResolvedRef::Int(_) => SapTy::Int, + TyResolvedRef::Bigint(_) => SapTy::Bigint, + TyResolvedRef::Float(_) => SapTy::Float, + TyResolvedRef::String(_) => SapTy::String, + TyResolvedRef::Bool(_) => SapTy::Bool, + TyResolvedRef::Null(_) => SapTy::Null, + TyResolvedRef::Media(media) => SapTy::Media(media.to_baml_media_kind()), + TyResolvedRef::LiteralInt(v) => { + SapTy::Literal(baml_type::Literal::Int(v.0), baml_type::Freshness::Regular) + } + TyResolvedRef::LiteralBigint(v) => SapTy::Literal( + baml_type::Literal::Bigint(v.0.clone()), + baml_type::Freshness::Regular, + ), + TyResolvedRef::LiteralString(v) => SapTy::Literal( + baml_type::Literal::String(v.0.to_string()), + baml_type::Freshness::Regular, + ), + TyResolvedRef::LiteralBool(v) => { + SapTy::Literal(baml_type::Literal::Bool(v.0), baml_type::Freshness::Regular) + } + TyResolvedRef::Array(a) => a.to_baml_ty(db), + TyResolvedRef::Map(m) => m.to_baml_ty(db), + TyResolvedRef::Class(c) => c.to_baml_ty(db), + TyResolvedRef::Enum(e) => e.to_baml_ty(db), + TyResolvedRef::EnumVariant(ev) => ev.to_baml_ty(db), + TyResolvedRef::Union(u) => u.to_baml_ty(db), + TyResolvedRef::StreamState(s) => s.to_baml_ty(db), + } + } +} + +impl ToBamlTy for Ty<'_, DefKey> { + fn to_baml_ty(&self, db: &TypeRefDb<'_, DefKey>) -> SapTy { + match self { + Ty::Resolved(resolved) => resolved.as_ref().to_baml_ty(db), + Ty::ResolvedRef(resolved_ref) => resolved_ref.to_baml_ty(db), + Ty::Unresolved(name) => match db.resolve_name(name) { + Some(TyResolvedRef::Class(class)) if class.name == *name => class.to_baml_ty(db), + Some(TyResolvedRef::Enum(enm)) if enm.name == *name => enm.to_baml_ty(db), + // Any other named entry is an alias. Keep it nominal here + // instead of recursively expanding aliases such as `json`. + Some(_) => SapTy::TypeAlias(name.clone()), + None => SapTy::unknown(), + }, + } + } +} + +impl ToBamlTy for ArrayTy<'_, DefKey> { + fn to_baml_ty(&self, db: &TypeRefDb<'_, DefKey>) -> SapTy { + let inner = self.ty.to_baml_ty(db); + SapTy::List(Box::new(inner)) + } +} + +impl ToBamlTy for MapTy<'_, DefKey> { + fn to_baml_ty(&self, db: &TypeRefDb<'_, DefKey>) -> SapTy { + SapTy::Map { + key: Box::new(self.key.to_baml_ty(db)), + value: Box::new(self.value.to_baml_ty(db)), + } + } +} + +impl ToBamlTy for ClassTy<'_, DefKey> { + fn to_baml_ty(&self, _db: &TypeRefDb<'_, DefKey>) -> SapTy { + SapTy::Class(self.name.clone(), Box::new([])) + } +} + +impl ToBamlTy for EnumTy<'_, DefKey> { + fn to_baml_ty(&self, _db: &TypeRefDb<'_, DefKey>) -> SapTy { + SapTy::Enum(self.name.clone()) + } +} + +impl ToBamlTy for EnumVariantTy<'_, DefKey> { + /// Loses variant specificity — maps back to the parent enum type. + fn to_baml_ty(&self, _db: &TypeRefDb<'_, DefKey>) -> SapTy { + SapTy::Enum(self.name.clone()) + } +} + +impl ToBamlTy for UnionTy<'_, DefKey> { + fn to_baml_ty(&self, db: &TypeRefDb<'_, DefKey>) -> SapTy { + let members: Vec = self.variants.iter().map(|v| v.to_baml_ty(db)).collect(); + SapTy::Union(members.into()) + } +} + +impl ToBamlTy for StreamStateTy<'_, DefKey> { + /// `StreamState` is a value-level concept; at the type level we return the inner type. + fn to_baml_ty(&self, db: &TypeRefDb<'_, DefKey>) -> SapTy { + self.value.to_baml_ty(db) + } +} + +impl MediaTy { + fn to_baml_media_kind(self) -> baml_type::MediaKind { + match self { + MediaTy::Image => baml_type::MediaKind::Image, + MediaTy::Audio => baml_type::MediaKind::Audio, + MediaTy::Pdf => baml_type::MediaKind::Pdf, + MediaTy::Video => baml_type::MediaKind::Video, + } + } +} diff --git a/baml_language/crates/bex_sap/src/to_external.rs b/baml_language/crates/bex_sap/src/to_external.rs deleted file mode 100644 index e914efc9a35..00000000000 --- a/baml_language/crates/bex_sap/src/to_external.rs +++ /dev/null @@ -1,275 +0,0 @@ -//! Conversion from SAP types and values to `BexExternalValue` and `SapTy`. -//! -//! This module provides the bridge between SAP's internal type/value representation -//! and the external value tree used by the engine to push results back onto the VM heap. - -use bex_external_types::BexExternalValue; -use indexmap::IndexMap; -use sys_types::{DefKey, SapTy}; - -use crate::{ - baml_value::{BamlStreamState, BamlValue}, - deserializer::types::BamlValueWithFlags, - sap_model::{ - ArrayTy, ClassTy, EnumTy, EnumVariantTy, MapTy, MediaTy, StreamStateTy, Ty, TyResolvedRef, - TypeRefDb, UnionTy, - }, -}; - -// ============================================================================ -// Type conversion: SAP types → SapTy -// ============================================================================ - -/// Convert a SAP type back to a `SapTy`. -/// -/// # Known simplifications (not round-trip safe) -/// -/// - `Optional(T)` was converted to `Union([Null, T])` on the way in; it stays as `Union` here. -/// - Aliases already resolved into the target type have no nominal name to recover; named alias -/// references inside containers are preserved. -/// - `EnumVariant` becomes `Enum` (variant specificity lost at the type level). -pub trait ToBamlTy { - fn to_baml_ty(&self, db: &TypeRefDb<'_, DefKey>) -> SapTy; -} - -impl ToBamlTy for TyResolvedRef<'_, DefKey> { - fn to_baml_ty(&self, db: &TypeRefDb<'_, DefKey>) -> SapTy { - match self { - TyResolvedRef::Int(_) => SapTy::Int, - TyResolvedRef::Bigint(_) => SapTy::Bigint, - TyResolvedRef::Float(_) => SapTy::Float, - TyResolvedRef::String(_) => SapTy::String, - TyResolvedRef::Bool(_) => SapTy::Bool, - TyResolvedRef::Null(_) => SapTy::Null, - TyResolvedRef::Media(media) => SapTy::Media(media.to_baml_media_kind()), - TyResolvedRef::LiteralInt(v) => { - SapTy::Literal(baml_type::Literal::Int(v.0), baml_type::Freshness::Regular) - } - TyResolvedRef::LiteralBigint(v) => SapTy::Literal( - baml_type::Literal::Bigint(v.0.clone()), - baml_type::Freshness::Regular, - ), - TyResolvedRef::LiteralString(v) => SapTy::Literal( - baml_type::Literal::String(v.0.to_string()), - baml_type::Freshness::Regular, - ), - TyResolvedRef::LiteralBool(v) => { - SapTy::Literal(baml_type::Literal::Bool(v.0), baml_type::Freshness::Regular) - } - TyResolvedRef::Array(a) => a.to_baml_ty(db), - TyResolvedRef::Map(m) => m.to_baml_ty(db), - TyResolvedRef::Class(c) => c.to_baml_ty(db), - TyResolvedRef::Enum(e) => e.to_baml_ty(db), - TyResolvedRef::EnumVariant(ev) => ev.to_baml_ty(db), - TyResolvedRef::Union(u) => u.to_baml_ty(db), - TyResolvedRef::StreamState(s) => s.to_baml_ty(db), - } - } -} - -impl ToBamlTy for Ty<'_, DefKey> { - fn to_baml_ty(&self, db: &TypeRefDb<'_, DefKey>) -> SapTy { - match self { - Ty::Resolved(resolved) => resolved.as_ref().to_baml_ty(db), - Ty::ResolvedRef(resolved_ref) => resolved_ref.to_baml_ty(db), - Ty::Unresolved(name) => match db.resolve_name(name) { - Some(TyResolvedRef::Class(class)) if class.name == *name => class.to_baml_ty(db), - Some(TyResolvedRef::Enum(enm)) if enm.name == *name => enm.to_baml_ty(db), - // Any other named entry is an alias. Keep it nominal here - // instead of recursively expanding aliases such as `json`. - Some(_) => SapTy::TypeAlias(name.clone()), - None => SapTy::unknown(), - }, - } - } -} - -impl ToBamlTy for ArrayTy<'_, DefKey> { - fn to_baml_ty(&self, db: &TypeRefDb<'_, DefKey>) -> SapTy { - let inner = self.ty.to_baml_ty(db); - SapTy::List(Box::new(inner)) - } -} - -impl ToBamlTy for MapTy<'_, DefKey> { - fn to_baml_ty(&self, db: &TypeRefDb<'_, DefKey>) -> SapTy { - SapTy::Map { - key: Box::new(self.key.to_baml_ty(db)), - value: Box::new(self.value.to_baml_ty(db)), - } - } -} - -impl ToBamlTy for ClassTy<'_, DefKey> { - fn to_baml_ty(&self, _db: &TypeRefDb<'_, DefKey>) -> SapTy { - SapTy::Class(self.name.clone(), Box::new([])) - } -} - -impl ToBamlTy for EnumTy<'_, DefKey> { - fn to_baml_ty(&self, _db: &TypeRefDb<'_, DefKey>) -> SapTy { - SapTy::Enum(self.name.clone()) - } -} - -impl ToBamlTy for EnumVariantTy<'_, DefKey> { - /// Loses variant specificity — maps back to the parent enum type. - fn to_baml_ty(&self, _db: &TypeRefDb<'_, DefKey>) -> SapTy { - SapTy::Enum(self.name.clone()) - } -} - -impl ToBamlTy for UnionTy<'_, DefKey> { - fn to_baml_ty(&self, db: &TypeRefDb<'_, DefKey>) -> SapTy { - let members: Vec = self.variants.iter().map(|v| v.to_baml_ty(db)).collect(); - SapTy::Union(members.into()) - } -} - -impl ToBamlTy for StreamStateTy<'_, DefKey> { - /// `StreamState` is a value-level concept; at the type level we return the inner type. - fn to_baml_ty(&self, db: &TypeRefDb<'_, DefKey>) -> SapTy { - self.value.to_baml_ty(db) - } -} - -impl MediaTy { - fn to_baml_media_kind(self) -> baml_type::MediaKind { - match self { - MediaTy::Image => baml_type::MediaKind::Image, - MediaTy::Audio => baml_type::MediaKind::Audio, - MediaTy::Pdf => baml_type::MediaKind::Pdf, - MediaTy::Video => baml_type::MediaKind::Video, - } - } -} - -// ============================================================================ -// Value conversion: SAP BamlValueWithFlags → BexExternalValue -// ============================================================================ - -/// Convert a lane type to the name-headed form `BexExternalValue` carries. -/// -/// The value metadata a parsed result carries has nowhere to put a declaration -/// identity, so it converts here — through the same per-call overlay spelling -/// the definition tables for this call were built with, so a returned value's -/// type names agree with the ones that call rendered. An anonymous declaration -/// has no resolvable name; its overlay spelling is meaningful only inside this -/// call and must not be treated as a lookup key beyond it. -fn to_value_metadata_ty(ty: &SapTy) -> baml_type::RuntimeTy { - match ty.try_map_heads(&mut |head: &DefKey| { - Ok::<_, std::convert::Infallible>(head.name().overlay_name()) - }) { - Ok(converted) => converted, - Err(never) => match never {}, - } -} - -/// Convert a SAP `BamlValueWithFlags` to a `BexExternalValue`, using the -/// type metadata from the deserializer to populate type information on -/// composite values (arrays, maps). -/// -/// Deserializer metadata (flags, scores) is discarded. -pub fn baml_value_to_external( - value: &BamlValueWithFlags<'_, '_, '_, DefKey>, - db: &TypeRefDb<'_, DefKey>, -) -> BexExternalValue { - baml_value_inner_to_external(&value.value, value.meta.ty, db) -} - -fn baml_value_inner_to_external( - value: &BamlValue<'_, '_, '_, DefKey>, - ty: TyResolvedRef<'_, DefKey>, - db: &TypeRefDb<'_, DefKey>, -) -> BexExternalValue { - match value { - BamlValue::String(s) => BexExternalValue::String(bex_str::BexStr::from(s.value.as_ref())), - BamlValue::Int(i) => BexExternalValue::Int(i.value), - BamlValue::Bigint(b) => BexExternalValue::Bigint(b.value.clone()), - BamlValue::Float(f) => BexExternalValue::Float(f.value), - BamlValue::Bool(b) => BexExternalValue::Bool(b.value), - BamlValue::Null(_) => BexExternalValue::Null, - BamlValue::Media(_) => { - unimplemented!("Media value conversion to BexExternalValue is not yet implemented") - } - BamlValue::Array(arr) => { - let element_type = to_value_metadata_ty(&match ty { - TyResolvedRef::Array(a) => a.ty.to_baml_ty(db), - _ => SapTy::unknown(), - }); - let items: Vec = arr - .value - .iter() - .map(|item| baml_value_to_external(item, db)) - .collect(); - BexExternalValue::Array { - element_type, - items, - } - } - BamlValue::Map(map) => { - let (key_type, value_type) = match ty { - TyResolvedRef::Map(m) => (m.key.to_baml_ty(db), m.value.to_baml_ty(db)), - _ => (SapTy::string(), SapTy::unknown()), - }; - let (key_type, value_type) = ( - to_value_metadata_ty(&key_type), - to_value_metadata_ty(&value_type), - ); - let entries: IndexMap = map - .value - .iter() - .map(|(k, v)| (k.to_string(), baml_value_to_external(v, db))) - .collect(); - BexExternalValue::Map { - key_type, - value_type, - entries, - } - } - BamlValue::Enum(e) => BexExternalValue::Variant { - // As for classes above: the qualified overlay spelling. - enum_name: e.name.name().overlay_name().to_string(), - variant_name: e.value.to_string(), - }, - BamlValue::Class(c) => { - let fields: IndexMap = c - .value - .iter() - .map(|(k, v)| (k.to_string(), baml_value_to_external(v, db))) - .collect(); - BexExternalValue::Instance { - // The landing resolves this against the per-call handle view, - // which is keyed by the *qualified* overlay spelling — so emit - // that, not the display name. One declaration has several name - // forms; the two sides of this hand-off must pick the same one. - class_name: c.name.name().overlay_name().to_string(), - type_args: vec![], - fields, - } - } - BamlValue::StreamState(state) => stream_state_to_external(state, db), - } -} - -fn stream_state_to_external( - state: &BamlStreamState<'_, '_, '_, DefKey>, - db: &TypeRefDb<'_, DefKey>, -) -> BexExternalValue { - let (inner, state_name) = match state { - BamlStreamState::Pending(v) => (v, "Pending"), - BamlStreamState::Incomplete(v) => (v, "Incomplete"), - BamlStreamState::Complete(v) => (v, "Complete"), - }; - let mut fields = IndexMap::new(); - fields.insert("value".to_string(), baml_value_to_external(inner, db)); - fields.insert( - "state".to_string(), - BexExternalValue::String(bex_str::BexStr::from(state_name)), - ); - BexExternalValue::Instance { - class_name: "StreamState".to_string(), - type_args: vec![], - fields, - } -} diff --git a/baml_language/crates/bex_vm/Cargo.toml b/baml_language/crates/bex_vm/Cargo.toml index 578f63f1d02..3195c19c78a 100644 --- a/baml_language/crates/bex_vm/Cargo.toml +++ b/baml_language/crates/bex_vm/Cargo.toml @@ -24,6 +24,7 @@ baml_compiler_diagnostics = { workspace = true } baml_linker_types = { workspace = true } baml_type = { workspace = true } bex_heap = { workspace = true } +bex_sap = { workspace = true } bex_str = { workspace = true } bex_vm_types = { workspace = true } btel_clock = { workspace = true } diff --git a/baml_language/crates/bex_vm/src/definitions.rs b/baml_language/crates/bex_vm/src/definitions.rs new file mode 100644 index 00000000000..2636b9af0ab --- /dev/null +++ b/baml_language/crates/bex_vm/src/definitions.rs @@ -0,0 +1,181 @@ +//! Class and enum definitions ([`sys_types::ClassDefinition`], +//! [`sys_types::EnumDefinition`]) read out of heap declarations. +//! +//! The one reader behind every definition table: the engine's program-wide +//! tables, its per-call runtime-type and runtime-package overlays, and +//! schema-aligned parsing. A class keeps every field with its `skip` flag, and +//! each consumer (output-format rendering, the JSON schema, the parser) leaves +//! skipped fields out itself. An enum leaves out skipped variants, which no +//! consumer reads. + +use bex_vm_types::{Class, Enum, RuntimeTy, TyTemplate, TypeHead}; +use sys_types::{ + ClassDefinition, ClassFieldDefinition, EnumDefinition, EnumVariantDefinition, SapTy, + SapTyTemplate, +}; + +/// The definition of `class`, every field included. +#[must_use] +pub fn class_definition(class: &Class) -> ClassDefinition { + ClassDefinition { + name: class.name.display_name().to_string(), + description: class.description.clone(), + docstring: class.docstring.clone(), + alias: class.alias.clone(), + stream_done: class.stream_done, + fields: class + .fields + .iter() + .map(|field| ClassFieldDefinition { + name: field.name.clone(), + field_type: lane_ty(&field.field_type), + field_template: Some(lane_template(&field.field_template)), + description: field.description.clone(), + docstring: field.docstring.clone(), + alias: field.alias.clone(), + skip: field.skip, + stream_done: field.stream_done, + must_exist: field.must_exist, + }) + .collect(), + } +} + +/// The definition of `enm`, without its skipped variants. +#[must_use] +pub fn enum_definition(enm: &Enum) -> EnumDefinition { + EnumDefinition { + name: enm.name.display_name().to_string(), + description: enm.description.clone(), + docstring: enm.docstring.clone(), + alias: enm.alias.clone(), + variants: enm + .variants + .iter() + .filter(|variant| !variant.skip) + .map(|variant| EnumVariantDefinition { + name: variant.name.clone(), + description: variant.description.clone(), + docstring: variant.docstring.clone(), + alias: variant.alias.clone(), + }) + .collect(), + } +} + +/// A declaration's type with each head as the tagged name it identifies. +/// +/// Total by invariant: a live declaration's heads are resolved pointers to +/// declarations, so each yields its identity and its own name. +#[must_use] +pub fn lane_ty(ty: &RuntimeTy) -> SapTy { + ty.try_map_heads(&mut TypeHead::to_tagged_name) + .unwrap_or_else(|head| { + unreachable!("a live declaration's field names a declaration: {head}") + }) +} + +/// [`lane_ty`] for a field's symbolic template, by the same invariant. +#[must_use] +pub fn lane_template(template: &TyTemplate) -> SapTyTemplate { + template + .try_map_heads(&mut TypeHead::to_tagged_name) + .unwrap_or_else(|head| { + unreachable!("a live declaration's field template names a declaration: {head}") + }) +} + +#[cfg(test)] +mod tests { + use baml_type::{DeclarationName, Name, typetag::TypeTag}; + use bex_vm_types::{ClassField, EnumVariant, types::Owner}; + use indexmap::IndexMap; + + use super::*; + + fn field( + name: &str, + field_type: RuntimeTy, + field_template: TyTemplate, + skip: bool, + ) -> ClassField { + ClassField { + name: name.to_string(), + field_type, + field_template, + description: None, + alias: None, + docstring: None, + other: IndexMap::new(), + skip, + stream_done: false, + must_exist: false, + runtime_type: None, + } + } + + fn variant(name: &str, skip: bool) -> EnumVariant { + EnumVariant { + name: name.to_string(), + description: None, + alias: None, + docstring: None, + other: IndexMap::new(), + skip, + } + } + + #[test] + fn class_definition_keeps_skipped_fields_flagged() { + let class = Class { + name: DeclarationName::Anonymous(Name::new("Person")), + fields: vec![ + field("name", RuntimeTy::String, TyTemplate::String, false), + field("age", RuntimeTy::Int, TyTemplate::Int, true), + ], + description: None, + alias: None, + docstring: None, + other: IndexMap::new(), + stream_done: false, + type_tag: TypeTag::fresh_dynamic(), + has_cleanup: false, + generic_param_count: 0, + owner: Owner::anonymous(), + methods: IndexMap::new(), + }; + + let definition = class_definition(&class); + + let fields: Vec<(&str, bool)> = definition + .fields + .iter() + .map(|field| (field.name.as_str(), field.skip)) + .collect(); + assert_eq!(fields, [("name", false), ("age", true)]); + assert_eq!(definition.fields[1].field_type, SapTy::Int); + } + + #[test] + fn enum_definition_leaves_out_skipped_variants() { + let enm = Enum { + name: DeclarationName::Anonymous(Name::new("Color")), + type_tag: TypeTag::fresh_dynamic(), + variants: vec![variant("Red", false), variant("Hidden", true)], + description: None, + alias: None, + docstring: None, + other: IndexMap::new(), + owner: Owner::anonymous(), + }; + + let definition = enum_definition(&enm); + + let names: Vec<&str> = definition + .variants + .iter() + .map(|variant| variant.name.as_str()) + .collect(); + assert_eq!(names, ["Red"]); + } +} diff --git a/baml_language/crates/bex_vm/src/lib.rs b/baml_language/crates/bex_vm/src/lib.rs index b14ed31d89d..6602b630454 100644 --- a/baml_language/crates/bex_vm/src/lib.rs +++ b/baml_language/crates/bex_vm/src/lib.rs @@ -14,6 +14,7 @@ pub(crate) mod array_index; mod call_specialize; pub mod compile_artifact; pub mod debug; +pub mod definitions; pub mod errors; pub mod indexable; pub mod kperf; diff --git a/baml_language/crates/bex_vm/src/package_baml/csv.rs b/baml_language/crates/bex_vm/src/package_baml/csv.rs index 2b31a2dbbb8..89ab0a6d292 100644 --- a/baml_language/crates/bex_vm/src/package_baml/csv.rs +++ b/baml_language/crates/bex_vm/src/package_baml/csv.rs @@ -190,34 +190,6 @@ fn done_value(vm: &mut BexVm) -> Result { // Handle-state access // ============================================================================= -/// Clone the `Arc` out of a `$rust_type` field so the caller can lock it while -/// still holding `&mut BexVm` for allocation. -fn state_arc( - vm: &BexVm, - holder: Value, - field_idx: usize, -) -> Result, VmRustFnError> { - let inst = vm.as_instance(&holder)?; - let handle = inst.load_field(field_idx); - let ptr = handle - .as_object_ptr() - .ok_or_else(|| VmInternalError::MissingNativeFunction { - name: "csv handle field is not an object".to_string(), - })?; - match vm.get_object(ptr) { - Object::RustData(arc) => arc.clone().downcast::().map_err(|_| { - VmRustFnError::InternalError(VmInternalError::MissingNativeFunction { - name: "csv handle holds an unexpected Rust type".to_string(), - }) - }), - _ => Err(VmRustFnError::InternalError( - VmInternalError::MissingNativeFunction { - name: "csv handle field is not RustData".to_string(), - }, - )), - } -} - fn lock(m: &Mutex) -> MutexGuard<'_, T> { m.lock().unwrap_or_else(std::sync::PoisonError::into_inner) } @@ -1490,7 +1462,7 @@ impl From for DecodeFail { } fn record_arc(vm: &BexVm, rec: Value) -> Result, VmRustFnError> { - state_arc::(vm, rec, 0) + Ok(vm.rust_data_field::(&rec, 0)?) } /// Decode a whole record into an instance of class `ty`. @@ -2088,7 +2060,7 @@ fn render_markdown(headers: &[String], rows: &[Vec], total_rows: usize) impl BamlClassCsvReader for PackageBamlImpl { fn skipped(vm: &mut BexVm, csvreader: &Value) -> Vec { - let Ok(st) = state_arc::>(vm, *csvreader, 0) else { + let Ok(st) = vm.rust_data_field::>(csvreader, 0) else { return Vec::new(); }; let infos: Vec = lock(&st).skipped.clone(); @@ -2103,7 +2075,7 @@ impl BamlClassCsvReader for PackageBamlImpl { } fn position(vm: &mut BexVm, csvreader: &Value) -> Value { - match state_arc::>(vm, *csvreader, 0) { + match vm.rust_data_field::>(csvreader, 0) { Ok(st) => { let (byte, line, record) = { let s = lock(&st); @@ -2121,7 +2093,7 @@ impl BamlClassCsvReader for PackageBamlImpl { } fn _poll(vm: &mut BexVm, csvreader: &Value) -> Result { - let st = state_arc::>(vm, *csvreader, 0)?; + let st = vm.rust_data_field::>(csvreader, 0)?; let polled = { let mut s = lock(&st); poll_record(&mut s) @@ -2142,7 +2114,7 @@ impl BamlClassCsvReader for PackageBamlImpl { } fn _poll_headers(vm: &mut BexVm, csvreader: &Value) -> Result { - let st = state_arc::>(vm, *csvreader, 0)?; + let st = vm.rust_data_field::>(csvreader, 0)?; let polled = { let mut s = lock(&st); poll_headers(&mut s) @@ -2257,7 +2229,7 @@ impl BamlClassCsvWriter for PackageBamlImpl { } fn text(vm: &mut BexVm, csvwriter: &Value) -> Result { - let st = state_arc::>(vm, *csvwriter, 0)?; + let st = vm.rust_data_field::>(csvwriter, 0)?; let text = { let s = lock(&st); s.buffer.clone() @@ -2279,7 +2251,7 @@ impl BamlClassCsvWriter for PackageBamlImpl { csvwriter: &Value, record: &[Value], ) -> Result { - let st = state_arc::>(vm, *csvwriter, 0)?; + let st = vm.rust_data_field::>(csvwriter, 0)?; let mut s = lock(&st); if s.closed { drop(s); @@ -2314,7 +2286,7 @@ impl BamlClassCsvWriter for PackageBamlImpl { csvwriter: &Value, names: &[Value], ) -> Result { - let st = state_arc::>(vm, *csvwriter, 0)?; + let st = vm.rust_data_field::>(csvwriter, 0)?; let mut s = lock(&st); if s.closed { drop(s); @@ -2432,7 +2404,7 @@ impl BamlNamespaceCsv for PackageBamlImpl { } }; - let st = state_arc::>(vm, *r, 0)?; + let st = vm.rust_data_field::>(r, 0)?; let header = lock(&st).header.clone(); for cf in &class_fields { @@ -2480,7 +2452,7 @@ impl BamlNamespaceCsv for PackageBamlImpl { Err(DecodeFail::Fatal(e)) => Err(e), Err(DecodeFail::Info(info)) => { let skip = info.kind == Kind::Decode && { - let st = state_arc::>(vm, *r, 0)?; + let st = vm.rust_data_field::>(r, 0)?; let mut s = lock(&st); if s.opts.skip_on_error { s.register_skip(&info); @@ -2557,7 +2529,7 @@ impl BamlNamespaceCsv for PackageBamlImpl { w: &Value, row: &Value, ) -> Result { - let st = state_arc::>(vm, *w, 0)?; + let st = vm.rust_data_field::>(w, 0)?; let mut s = lock(&st); if s.closed { drop(s); @@ -2584,7 +2556,7 @@ impl BamlNamespaceCsv for PackageBamlImpl { w: &Value, rows: &[Value], ) -> Result { - let st = state_arc::>(vm, *w, 0)?; + let st = vm.rust_data_field::>(w, 0)?; let mut s = lock(&st); if s.closed { drop(s); diff --git a/baml_language/crates/bex_vm/src/package_baml/mod.rs b/baml_language/crates/bex_vm/src/package_baml/mod.rs index d003fffa7fe..f7bcdb3399d 100644 --- a/baml_language/crates/bex_vm/src/package_baml/mod.rs +++ b/baml_language/crates/bex_vm/src/package_baml/mod.rs @@ -44,6 +44,7 @@ mod regex; pub(crate) mod resolve; pub(crate) use resolve::ImplResolver; pub(crate) mod root; +mod sap; mod spawn; pub use spawn::{ SpawnLaunch, alloc_projected_cancel_token, cancel_token_members, plan_body, diff --git a/baml_language/crates/bex_vm/src/package_baml/regex.rs b/baml_language/crates/bex_vm/src/package_baml/regex.rs index 08bae5a897c..8729f467930 100644 --- a/baml_language/crates/bex_vm/src/package_baml/regex.rs +++ b/baml_language/crates/bex_vm/src/package_baml/regex.rs @@ -215,24 +215,7 @@ fn regex_value(vm: &mut BexVm, prog: Program) -> Value { /// Clone the `Arc` out of `Regex._handle` so the program stays alive while /// `&mut BexVm` is borrowed for allocation. fn program_of(vm: &BexVm, regex: Value) -> Result, VmRustFnError> { - let instance = vm.as_instance(®ex)?; - let ptr = instance.load_field(0).as_object_ptr().ok_or_else(|| { - VmInternalError::MissingNativeFunction { - name: "baml.regex.Regex._handle is not an object".to_string(), - } - })?; - match vm.get_object(ptr) { - Object::RustData(data) => data.clone().downcast::().map_err(|_| { - VmRustFnError::InternalError(VmInternalError::MissingNativeFunction { - name: "baml.regex.Regex._handle holds an unexpected Rust type".to_string(), - }) - }), - _ => Err(VmRustFnError::InternalError( - VmInternalError::MissingNativeFunction { - name: "baml.regex.Regex._handle is not RustData".to_string(), - }, - )), - } + Ok(vm.rust_data_field::(®ex, 0)?) } /// A search that ran out of budget is a panic, not a "no match". diff --git a/baml_language/crates/bex_vm/src/package_baml/sap.rs b/baml_language/crates/bex_vm/src/package_baml/sap.rs new file mode 100644 index 00000000000..172a307c53f --- /dev/null +++ b/baml_language/crates/bex_vm/src/package_baml/sap.rs @@ -0,0 +1,503 @@ +//! `baml.sap`: schema-aligned parsing (SAP) of an LLM's text into a declared +//! type. +//! +//! VM natives, not sys ops: parsing does no IO, so it runs on the calling VM +//! with the heap at hand. `_new_parse_cache` builds the parse model from the +//! declarations `T` reaches, and the parse methods allocate their result +//! directly on the heap, so no definition table is copied and no value tree is +//! landed by name. + +use std::{any::Any, collections::HashMap, sync::Arc}; + +use baml_type::{DeclarationName, Literal, typetag::TypeTag}; +use bex_heap::TlabHolder; +use bex_sap::{ + CompiledSapModel, + baml_value::BamlValue, + deserializer::{coercer::ParsingContext, types::BamlValueWithFlags}, + jsonish, + sap_model::{TyResolvedRef, TypeCtx, TypeRefDb}, + to_baml_ty::ToBamlTy, +}; +use bex_str::BexStr; +use bex_vm_types::{ + HeapPtr, RealizedTy, RuntimeTy, TypeHead, + types::{Instance, Object, Value}, +}; +use indexmap::IndexMap; +use sys_types::{DefKey, SapTy}; + +use super::{BamlNamespaceSap, BamlNamespaceSap_ParseCache, PackageBamlImpl}; +use crate::{ + BexVm, definitions, + errors::{VmBamlError, VmInternalError, VmPanic, VmRustFnError}, + reachable, +}; + +/// The `_data` of a `baml.sap._ParseCache`: the parse model for `T`. +/// +/// A `RustData` payload is a GC leaf, so the model holds no heap pointers: it +/// names declarations by tagged name, and each parse maps those names back to +/// the declarations through [`Declarations`]. +struct ParseCache { + model: CompiledSapModel, +} + +impl BamlNamespaceSap for PackageBamlImpl { + fn _new_parse_cache(vm: &mut BexVm) -> Result { + let target = parse_target(vm)?; + let model = build_model(vm, &target).map_err(|e| { + // `throws never`: a `T` the parser cannot model comes from the + // caller's own `parse`, a program bug rather than bad input. + VmPanic::UserPanic { + message: format!("schema-aligned parsing cannot model this type: {e}"), + } + })?; + let data: Arc = Arc::new(ParseCache { model }); + let data = Value::object(vm.alloc_rust_data(data)); + let class = vm.resolve_class("baml.sap._ParseCache"); + let cache = vm.tlab.alloc(Object::Instance(Instance::new( + class, + Box::new([target]), + vec![data], + ))); + Ok(Value::object(cache)) + } +} + +impl BamlNamespaceSap_ParseCache for PackageBamlImpl { + fn _parse_final(vm: &mut BexVm, cache: &Value, json: &BexStr) -> Result { + let parsed = parse(vm, *cache, json, Completion::Final)?; + parsed.ok_or_else(|| llm_client("SAP parse returned no value when complete")) + } + + fn _parse_partial( + vm: &mut BexVm, + cache: &Value, + json: &BexStr, + ) -> Result { + match parse(vm, *cache, json, Completion::Partial)? { + Some(value) => Ok(value), + None => { + let no_yield = vm.resolve_class("baml.sap._NoYield"); + Ok(Value::object(vm.alloc_instance(no_yield, Vec::new()))) + } + } + } +} + +/// Whether the text is the whole reply, or the part a stream has received. +#[derive(Clone, Copy, PartialEq, Eq)] +enum Completion { + Final, + Partial, +} + +/// The `T` of the current `_new_parse_cache` or `_ParseCache` method +/// call: the first type argument either way. +fn parse_target(vm: &BexVm) -> Result { + vm.current_call_type_args().first().cloned().ok_or_else(|| { + VmInternalError::SapValue { + message: "a `baml.sap` native was called without its `T`".to_string(), + } + .into() + }) +} + +/// The declarations the parse model of `target` needs: those `target` names, +/// then those a reached class's field types name (except `@skip` fields, which +/// the parser drops) and those a reached alias's body names. +fn parsed_declarations(vm: &BexVm, target: &RealizedTy) -> Vec { + reachable::declarations_reached(vm, target, |object, visit| match object { + Object::Class(class) => { + for field in class.fields.iter().filter(|field| !field.skip) { + field.field_type.visit_heads(&mut |head| visit(head)); + } + } + Object::TypeAlias(alias) => alias.definition.visit_heads(&mut |head| visit(head)), + _ => {} + }) +} + +/// The parse model for `target`, built from the declarations it reaches. +fn build_model( + vm: &BexVm, + target: &RealizedTy, +) -> Result { + let mut classes = IndexMap::new(); + let mut enums = IndexMap::new(); + let mut aliases = HashMap::new(); + for ptr in parsed_declarations(vm, target) { + match vm.get_object(ptr) { + Object::Class(class) => { + classes.insert( + DefKey::new(class.type_tag, class.name.clone()), + definitions::class_definition(class), + ); + } + Object::Enum(enm) => { + enums.insert( + DefKey::new(enm.type_tag, enm.name.clone()), + definitions::enum_definition(enm), + ); + } + Object::TypeAlias(alias) => { + aliases.insert( + DefKey::new( + alias.type_tag, + DeclarationName::Declared(alias.name.clone()), + ), + definitions::lane_ty(&RuntimeTy::from(alias.definition.clone())), + ); + } + _ => {} + } + } + let target = definitions::lane_ty(&RuntimeTy::from(target.clone())); + CompiledSapModel::from_type_ctx(TypeCtx::new(&classes, enums, &aliases), target) +} + +/// Parse `text` with the model in `cache` and allocate the result, or `None` +/// for a partial reply that does not parse yet. +fn parse( + vm: &mut BexVm, + cache: Value, + text: &BexStr, + completion: Completion, +) -> Result, VmRustFnError> { + let cache = vm.rust_data_field::(&cache, 0)?; + let target_type = parse_target(vm)?; + let model = &cache.model; + // An unresolvable target is a schema bug, not an incomplete reply. + let target = model + .resolved_target() + .map_err(|e| llm_client(e.to_string()))?; + let is_done = completion == Completion::Final; + let value = match jsonish::parse(text.as_str(), jsonish::ParseOptions::default(), is_done) { + Ok(value) => value, + Err(e) if is_done => return Err(llm_client(e.to_string())), + // A reply that opens with a fence or with prose before the JSON cannot + // be coerced until the JSON starts: skip the partial, wait for text. + Err(_) => return Ok(None), + }; + let context = ParsingContext::new(model.db()); + let parsed = match TyResolvedRef::coerce(&context, target, &value) { + Ok(parsed) => parsed, + Err(e) if is_done => return Err(llm_client(e.to_string())), + Err(_) => return Ok(None), + }; + // Only a value to allocate needs the declarations: a partial that does not + // parse yet skips the walk. + parsed + .map(|parsed| { + let declarations = Declarations::reached_by(vm, &target_type); + to_heap(vm, &parsed, model.db(), &declarations) + }) + .transpose() +} + +/// A parse failure, as `baml.errors.LlmClient`. +fn llm_client(message: impl Into) -> VmRustFnError { + VmBamlError::LlmClient { + message: message.into(), + } + .into() +} + +/// The declarations a parse target reaches, by tag, for mapping the model's +/// tagged names back to the heap. Built for one call: a parse never yields, so +/// no collection can move them meanwhile. +struct Declarations(HashMap); + +impl Declarations { + fn reached_by(vm: &BexVm, target: &RealizedTy) -> Self { + Self( + parsed_declarations(vm, target) + .into_iter() + .filter_map(|ptr| { + let tag = match vm.get_object(ptr) { + Object::Class(class) => class.type_tag, + Object::Enum(enm) => enm.type_tag, + Object::TypeAlias(alias) => alias.type_tag, + _ => return None, + }; + Some((tag, ptr)) + }) + .collect(), + ) + } + + fn get(&self, name: &DefKey) -> Result { + self.0.get(&name.tag()).copied().ok_or_else(|| { + invariant(format!( + "the parsed value names `{name}`, which its target does not reach" + )) + }) + } + + /// `ty` with each tagged name replaced by the head of its declaration. + fn realize(&self, ty: &SapTy) -> Result { + let ty = ty.try_map_heads(&mut |name: &DefKey| { + self.get(name).map(|ptr| TypeHead::new(ptr, name.tag())) + })?; + RealizedTy::try_from(ty) + .map_err(|e| invariant(format!("a parsed value's type is not realized: {e}"))) + } +} + +fn invariant(message: String) -> VmRustFnError { + VmInternalError::SapValue { message }.into() +} + +/// Allocate `value`: the same value the parser produced, on the heap. +fn to_heap( + vm: &mut BexVm, + value: &BamlValueWithFlags<'_, '_, '_, DefKey>, + db: &TypeRefDb<'_, DefKey>, + declarations: &Declarations, +) -> Result { + Ok(match &value.value { + BamlValue::String(s) => Value::object(vm.alloc_string(s.value.as_ref())), + BamlValue::Int(i) => Value::try_int(i.value).ok_or_else(|| { + llm_client(format!( + "integer {} is outside the BAML integer range [{}, {}]", + i.value, + Value::INT_MIN, + Value::INT_MAX + )) + })?, + BamlValue::Bigint(b) => vm.try_alloc_bigint(Arc::new(b.value.clone()))?, + BamlValue::Float(f) => Value::object(vm.alloc_float(f.value)), + BamlValue::Bool(b) => Value::bool(b.value), + BamlValue::Null(_) => Value::NULL, + BamlValue::Media(_) => { + return Err(invariant( + "the parser produced a media value, which it cannot parse".to_string(), + )); + } + BamlValue::Array(array) => { + let element = match value.meta.ty { + TyResolvedRef::Array(array) => array.ty.to_baml_ty(db), + _ => SapTy::unknown(), + }; + let element = declarations.realize(&element)?; + let items = array + .value + .iter() + .map(|item| to_heap(vm, item, db, declarations)) + .collect::, _>>()?; + Value::object(vm.alloc_array(element, items)) + } + BamlValue::Map(map) => { + let (key, value_ty) = match value.meta.ty { + TyResolvedRef::Map(map) => (map.key.to_baml_ty(db), map.value.to_baml_ty(db)), + _ => (SapTy::string(), SapTy::unknown()), + }; + let (key, value_ty) = ( + declarations.realize(&key)?, + declarations.realize(&value_ty)?, + ); + let mut entries = IndexMap::with_capacity(map.value.len()); + for (k, v) in &map.value { + entries.insert(BexStr::from(k.as_ref()), to_heap(vm, v, db, declarations)?); + } + Value::object(vm.alloc_map(key, value_ty, entries)) + } + BamlValue::Enum(e) => { + let enm = declarations.get(e.name)?; + let Object::Enum(declaration) = vm.get_object(enm) else { + return Err(invariant(format!("`{}` is not an enum", e.name))); + }; + let index = declaration + .variants + .iter() + .position(|variant| variant.name == e.value) + .ok_or_else(|| { + invariant(format!("enum `{}` has no variant `{}`", e.name, e.value)) + })?; + Value::object(vm.alloc_variant(enm, index)) + } + BamlValue::Class(c) => { + let class = declarations.get(c.name)?; + let Object::Class(declaration) = vm.get_object(class) else { + return Err(invariant(format!("`{}` is not a class", c.name))); + }; + let fields: Vec<(String, bool, RuntimeTy)> = declaration + .fields + .iter() + .map(|field| (field.name.clone(), field.skip, field.field_type.clone())) + .collect(); + let mut values = Vec::with_capacity(fields.len()); + for (name, skip, field_type) in fields { + let field_value = match c.value.get(name.as_str()) { + Some(parsed) => to_heap(vm, parsed, db, declarations)?, + None if skip => skipped_field_default(vm, c.name, &name, &field_type)?, + None => { + return Err(invariant(format!( + "the parsed `{}` has no field `{name}`", + c.name + ))); + } + }; + values.push(field_value); + } + Value::object(vm.alloc_instance(class, values)) + } + BamlValue::StreamState(_) => { + return Err(invariant( + "the parser produced a stream state, which no target declares".to_string(), + )); + } + }) +} + +/// The value a `@skip` field holds in a parsed instance. The parser never reads +/// the field, so it gets its type's empty value. +fn skipped_field_default( + vm: &mut BexVm, + class: &DefKey, + field: &str, + field_type: &RuntimeTy, +) -> Result { + let no_default = |ty: &str| -> VmRustFnError { + VmPanic::UserPanic { + message: format!( + "the `@skip` field `{class}.{field}` has type {ty}, which has no default value \ + for schema-aligned parsing to fill in" + ), + } + .into() + }; + let realized = RealizedTy::try_from(field_type.clone()) + .map_err(|_| no_default("of a generic parameter"))?; + empty_value(vm, &realized, &mut Vec::new())?.ok_or_else(|| no_default(&realized.to_string())) +} + +/// The empty value of `ty`: `null` for a nullable type, the zero of a number, +/// `false`, `""`, an empty list or map, a literal's own value, an enum's first +/// variant, and a class of its fields' empty values. `None` when `ty` has none +/// (media, functions, a class that contains itself, ...). +/// +/// `within` holds the classes and aliases being built, to stop a type that +/// needs its own empty value. +fn empty_value( + vm: &mut BexVm, + ty: &RealizedTy, + within: &mut Vec, +) -> Result, VmRustFnError> { + Ok(Some(match ty { + RealizedTy::Null | RealizedTy::Unknown => Value::NULL, + RealizedTy::Int => Value::int(0), + RealizedTy::Bigint => vm.try_alloc_bigint(Arc::new(num_bigint::BigInt::default()))?, + RealizedTy::Float => Value::object(vm.alloc_float(0.0)), + RealizedTy::Bool => Value::bool(false), + RealizedTy::String => Value::object(vm.alloc_string("")), + RealizedTy::Uint8Array => Value::object(vm.alloc_uint8array(Vec::new())), + RealizedTy::Literal(literal, _) => match literal { + Literal::Int(n) => Value::int(*n), + Literal::Bigint(n) => vm.try_alloc_bigint(Arc::new(n.clone()))?, + Literal::Float(text) => match text.parse::() { + Ok(f) => Value::object(vm.alloc_float(f)), + Err(_) => return Ok(None), + }, + Literal::String(s) => Value::object(vm.alloc_string(s.as_str())), + Literal::Bool(b) => Value::bool(*b), + }, + RealizedTy::List(element) => Value::object(vm.alloc_array((**element).clone(), Vec::new())), + RealizedTy::Map { key, value } => { + Value::object(vm.alloc_map((**key).clone(), (**value).clone(), IndexMap::new())) + } + RealizedTy::Enum(head) => { + let Object::Enum(enm) = vm.get_object(head.ptr()) else { + return Ok(None); + }; + let Some(index) = enm.variants.iter().position(|variant| !variant.skip) else { + return Ok(None); + }; + Value::object(vm.alloc_variant(head.ptr(), index)) + } + RealizedTy::EnumVariant(head, name) => { + let Object::Enum(enm) = vm.get_object(head.ptr()) else { + return Ok(None); + }; + let Some(index) = enm + .variants + .iter() + .position(|variant| variant.name == name.as_str()) + else { + return Ok(None); + }; + Value::object(vm.alloc_variant(head.ptr(), index)) + } + RealizedTy::Union(members) => { + if members + .iter() + .any(|member| matches!(member, RealizedTy::Null)) + { + Value::NULL + } else { + let mut found = None; + for member in members { + if let Some(value) = empty_value(vm, member, within)? { + found = Some(value); + break; + } + } + return Ok(found); + } + } + RealizedTy::Class(head, type_args) => { + let class = head.ptr(); + if within.contains(&class) { + return Ok(None); + } + let Object::Class(declaration) = vm.get_object(class) else { + return Ok(None); + }; + let templates: Vec<_> = declaration + .fields + .iter() + .map(|field| field.field_template.clone()) + .collect(); + within.push(class); + let mut values = Vec::with_capacity(templates.len()); + for template in &templates { + let field_ty = vm.realize_field_ty(template, type_args); + let Some(value) = empty_value(vm, &field_ty, within)? else { + within.pop(); + return Ok(None); + }; + values.push(value); + } + within.pop(); + Value::object( + vm.tlab + .alloc_instance_with_type_args(class, type_args.clone(), values), + ) + } + RealizedTy::TypeAlias(head) => { + let alias = head.ptr(); + if within.contains(&alias) { + return Ok(None); + } + let Object::TypeAlias(declaration) = vm.get_object(alias) else { + return Ok(None); + }; + let definition = declaration.definition.clone(); + within.push(alias); + let value = empty_value(vm, &definition, within)?; + within.pop(); + return Ok(value); + } + RealizedTy::Media(_) + | RealizedTy::Function { .. } + | RealizedTy::Future(..) + | RealizedTy::Interface(..) + | RealizedTy::RustType + | RealizedTy::Type + | RealizedTy::Resource + | RealizedTy::PromptAst + | RealizedTy::Void + | RealizedTy::Never => return Ok(None), + })) +} diff --git a/baml_language/crates/bex_vm/src/reachable.rs b/baml_language/crates/bex_vm/src/reachable.rs index 3f7752881f6..66430e59c1e 100644 --- a/baml_language/crates/bex_vm/src/reachable.rs +++ b/baml_language/crates/bex_vm/src/reachable.rs @@ -121,14 +121,17 @@ pub fn runtime_nominals_under_permit( (classes, enums) } -/// Every class and enum `ty` reaches, including compile-time declarations. -/// -/// Most runtime consumers intentionally stop at the immutable program image, -/// but source rendering must emit a standalone graph: a runtime class field -/// that names a static class needs that static declaration in the output too. +/// Every declaration `ty` reaches when each reached declaration leads on to the +/// heads `follow` reports for it, in the order a depth-first walk meets them: +/// the heads of `ty`, then each declaration's followed heads. #[must_use] -pub fn all_nominals(vm: &BexVm, ty: &bex_vm_types::RealizedTy) -> (Vec, Vec) { +pub fn declarations_reached( + vm: &BexVm, + ty: &bex_vm_types::RealizedTy, + follow: impl Fn(&Object, &mut dyn FnMut(&bex_vm_types::TypeHead)), +) -> Vec { let mut found = Vec::new(); + let mut seen = std::collections::HashSet::new(); let mut pending = Vec::new(); ty.visit_heads(&mut |head| { if head.is_resolved() { @@ -137,13 +140,12 @@ pub fn all_nominals(vm: &BexVm, ty: &bex_vm_types::RealizedTy) -> (Vec, }); pending.reverse(); while let Some(ptr) = pending.pop() { - if found.contains(&ptr) { + if !seen.insert(ptr) { continue; } found.push(ptr); let mut next = Vec::new(); - let object = vm.get_object(ptr); - bex_vm_types::head_walk::visit_object_heads(object, &mut |head| { + follow(vm.get_object(ptr), &mut |head| { if head.is_resolved() { next.push(head.ptr()); } @@ -151,10 +153,29 @@ pub fn all_nominals(vm: &BexVm, ty: &bex_vm_types::RealizedTy) -> (Vec, next.reverse(); pending.extend(next); } + found +} +/// Every declaration `ty` reaches (classes, enums, interfaces, type aliases), +/// including compile-time ones, through every head a declaration holds (field +/// types, templates, alias bodies, ...). +#[must_use] +pub fn all_declarations(vm: &BexVm, ty: &bex_vm_types::RealizedTy) -> Vec { + declarations_reached(vm, ty, |object, visit| { + bex_vm_types::head_walk::visit_object_heads(object, &mut |head| visit(head)); + }) +} + +/// Every class and enum `ty` reaches, including compile-time declarations. +/// +/// Most runtime consumers intentionally stop at the immutable program image, +/// but source rendering must emit a standalone graph: a runtime class field +/// that names a static class needs that static declaration in the output too. +#[must_use] +pub fn all_nominals(vm: &BexVm, ty: &bex_vm_types::RealizedTy) -> (Vec, Vec) { let mut classes = Vec::new(); let mut enums = Vec::new(); - for ptr in found { + for ptr in all_declarations(vm, ty) { match vm.get_object(ptr) { Object::Class(_) => classes.push(ptr), Object::Enum(_) => enums.push(ptr), diff --git a/baml_language/crates/bex_vm/src/vm.rs b/baml_language/crates/bex_vm/src/vm.rs index cb6003d1e82..8a02525584b 100644 --- a/baml_language/crates/bex_vm/src/vm.rs +++ b/baml_language/crates/bex_vm/src/vm.rs @@ -4940,6 +4940,38 @@ impl BexVm { } } + /// Clone the `Arc` out of the `$rust_type` field `field` of the instance + /// `holder`, so a native can keep using the payload while it borrows + /// `&mut BexVm` to allocate. + pub fn rust_data_field( + &self, + holder: &Value, + field: usize, + ) -> Result, VmInternalError> { + let handle = self.as_instance(holder)?.load_field(field); + let Some(ptr) = handle.as_object_ptr() else { + return Err(VmInternalError::TypeError { + expected: Type::Object(ObjectType::RustData), + got: self.type_of(&handle), + }); + }; + match self.get_object(ptr) { + Object::RustData(arc) => { + let got = arc.as_ref().type_id(); + arc.clone() + .downcast::() + .map_err(|_| VmInternalError::RustTypeError { + expected: TypeId::of::(), + got, + }) + } + _ => Err(VmInternalError::TypeError { + expected: Type::Object(ObjectType::RustData), + got: self.type_of(&handle), + }), + } + } + /// Extract an `&Instance` from a `Value` carrying a heap-object pointer. /// /// Used by generated glue code to construct `view::` structs. diff --git a/baml_language/crates/bex_vm_types/src/errors.rs b/baml_language/crates/bex_vm_types/src/errors.rs index 452c13a106d..b04b02188a2 100644 --- a/baml_language/crates/bex_vm_types/src/errors.rs +++ b/baml_language/crates/bex_vm_types/src/errors.rs @@ -244,6 +244,10 @@ pub enum VmInternalError { #[error("missing native function: {name}")] MissingNativeFunction { name: String }, + /// Schema-aligned parsing produced a value its target's declarations + /// cannot hold: the parse model and the heap disagree. + #[error("schema-aligned parsing produced a value its target cannot hold: {message}")] + SapValue { message: String }, #[error("unexpected constant kind: expected a TyTemplate constant at this index")] UnexpectedConstantKind, diff --git a/baml_language/crates/sys_ops/Cargo.toml b/baml_language/crates/sys_ops/Cargo.toml index 93e2d6f7c8a..a2e5668183f 100644 --- a/baml_language/crates/sys_ops/Cargo.toml +++ b/baml_language/crates/sys_ops/Cargo.toml @@ -25,7 +25,6 @@ baml_type = { workspace = true } bex_external_types = { workspace = true } bex_heap = { workspace = true } bex_resource_types = { workspace = true } -bex_sap = { workspace = true } bex_vm_types = { workspace = true } sys_types = { workspace = true } indexmap = { workspace = true } diff --git a/baml_language/crates/sys_ops/src/lib.rs b/baml_language/crates/sys_ops/src/lib.rs index 022c7cc3d50..553839ed03c 100644 --- a/baml_language/crates/sys_ops/src/lib.rs +++ b/baml_language/crates/sys_ops/src/lib.rs @@ -94,7 +94,6 @@ pub use io_adapter::build_runtime_io; // ============================================================================ pub mod output_format; -pub mod sap; // ============================================================================ // Blanket IO LLM implementation @@ -150,9 +149,8 @@ fn llm_function_lookup_error( } } -/// Blanket impl — schema-aligned parsing, backing both public `baml.sap.parse` -/// and incremental `ai.stream.Stream` parsing. -/// +/// Blanket impl — `ai.OutputFormat`: rendering the output-format schema with +/// the caller's options. impl io::IoClassAiOutputFormat for T { #[allow(clippy::too_many_arguments)] fn _render( @@ -899,85 +897,6 @@ pub fn get_return_type_op( SysOpOutput::ok(info.return_type.clone()) } -/// Blanket impl — schema-aligned parsing, backing both public `baml.sap.parse` -/// and incremental `ai.stream.Stream` parsing. All three are free functions -/// (see `ns_sap/sap.baml`), so each carries its own `T` type-arg operand; the -/// cache already holds the compiled model, so the two parse entry points -/// ignore theirs. -impl io::IoClassSapParseCache for T { - fn _parse_final( - &self, - _heap: &std::sync::Arc, - _call_id: CallId, - cache: io::owned::sap::ParseCache, - json: String, - _type_arg_0: ::sys_types::SapTy, - ctx: &SysOpContext, - ) -> SysOpOutput { - let Ok(sap) = cache._data.clone().downcast::() else { - return SysOpOutput::err(VmInternalError::RustTypeError { - expected: std::any::TypeId::of::(), - got: cache._data.type_id(), - }); - }; - SysOpOutput::Ready( - crate::sap::execute_sap_parse_final(&json, &sap, ctx).map_err(VmRustFnError::from), - ) - } - - fn _parse_partial( - &self, - _heap: &std::sync::Arc, - _call_id: CallId, - cache: io::owned::sap::ParseCache, - json: String, - _type_arg_0: ::sys_types::SapTy, - ctx: &SysOpContext, - ) -> SysOpOutput { - let Ok(sap) = cache._data.clone().downcast::() else { - return SysOpOutput::err(VmInternalError::RustTypeError { - expected: std::any::TypeId::of::(), - got: cache._data.type_id(), - }); - }; - let result = match crate::sap::execute_sap_parse_partial(&json, &sap, ctx) { - Ok(Some(value)) => Ok(value), - Ok(None) => Ok(BexExternalValue::instance( - "baml.sap._NoYield", - ::indexmap::IndexMap::new(), - )), - Err(e) => Err(VmRustFnError::from(e)), - }; - SysOpOutput::Ready(result) - } -} - -impl io::IoNamespaceSap for T { - fn _new_parse_cache( - &self, - _heap: &std::sync::Arc, - _call_id: CallId, - target: ::sys_types::SapTy, - ctx: &SysOpContext, - ) -> SysOpOutput { - let compiled = match ::bex_sap::CompiledSapModel::from_sys_op_context(ctx, target) { - Ok(compiled) => compiled, - Err(e) => { - // `_new_parse_cache` declares `throws never`, and the type - // arguments that reach it come from the caller's own - // `parse` — a `T` schema-aligned parsing cannot model is - // a program bug, not a recoverable condition, so it panics. - return SysOpOutput::err(VmPanic::UserPanic { - message: format!("schema-aligned parsing cannot model this type: {e}"), - }); - } - }; - let sap = crate::sap::SapParseCache::new(compiled); - let data: std::sync::Arc = std::sync::Arc::new(sap); - SysOpOutput::ok(io::owned::sap::ParseCache { _data: data }) - } -} - /// Wrap an `OutputFormatContent` into the generated `owned::ai::OutputFormat` handle. fn wrap_output_format( content: std::sync::Arc, diff --git a/baml_language/crates/sys_ops/src/sap.rs b/baml_language/crates/sys_ops/src/sap.rs deleted file mode 100644 index 03255b042c0..00000000000 --- a/baml_language/crates/sys_ops/src/sap.rs +++ /dev/null @@ -1,173 +0,0 @@ -use ::bex_sap::sap_model; -use ::sys_types::DefKey; - -/// Cached schema information for incremental or one-shot SAP parsing. -pub struct SapParseCache { - types: bex_sap::CompiledSapModel, -} - -impl SapParseCache { - pub fn new(types: bex_sap::CompiledSapModel) -> Self { - Self { types } - } - - pub fn db(&self) -> &sap_model::TypeRefDb<'_, DefKey> { - self.types.db() - } - - pub fn ty(&self) -> &sap_model::Ty<'_, DefKey> { - self.types.ty() - } - - pub fn ty_resolved(&self) -> Result, &DefKey> { - self.db().resolve(self.ty()) - } -} - -/// Errors that can occur during LLM operations. Relocated verbatim from -/// `sys_llm`; only `ParseResponseError`, `JsonishError` and `SapError` are -/// still constructible now that the SAP parse entry points are the sole users. -#[derive(Debug, thiserror::Error)] -pub enum LlmOpError { - #[error("Expected {expected}, got {actual}")] - TypeError { - expected: &'static str, - actual: String, - }, - - #[error("Parse response error: {0}")] - ParseResponseError(String), - - #[error("Jsonish error: {0}")] - JsonishError(::bex_sap::jsonish::JsonishError), - - #[error("SAP error: {0}")] - SapError(::bex_sap::deserializer::coercer::ParsingError), -} - -impl From for ::sys_types::VmRustFnError { - fn from(e: LlmOpError) -> Self { - let baml: ::sys_types::VmBamlError = match e { - LlmOpError::TypeError { expected, actual } => { - ::sys_types::VmBamlError::InvalidArgument { - message: format!("expected {expected}, got {actual}"), - } - } - LlmOpError::ParseResponseError(e) => ::sys_types::VmBamlError::LlmClient { message: e }, - LlmOpError::JsonishError(e) => ::sys_types::VmBamlError::LlmClient { - message: e.to_string(), - }, - LlmOpError::SapError(e) => ::sys_types::VmBamlError::LlmClient { - message: e.to_string(), - }, - }; - baml.into() - } -} - -pub fn execute_sap_parse_final( - json: &str, - sap: &SapParseCache, - _ctx: &::sys_types::SysOpContext, -) -> Result { - // === Jsonish === - let jsonish_options = ::bex_sap::jsonish::ParseOptions::default(); - let jsonish = - ::bex_sap::jsonish::parse(json, jsonish_options, true).map_err(LlmOpError::JsonishError)?; - - let parse_ctx = ::bex_sap::deserializer::coercer::ParsingContext::new(sap.db()); - let target = sap - .ty_resolved() - .map_err(|err| parse_ctx.error_type_resolution(err)) - .map_err(LlmOpError::SapError)?; - let parsed = ::bex_sap::sap_model::TyResolvedRef::coerce(&parse_ctx, target, &jsonish) - .map_err(LlmOpError::SapError)? - .ok_or_else(|| { - LlmOpError::ParseResponseError("SAP parse returned no value when complete".to_string()) - })?; - - // === Convert back to baml === - Ok(::bex_sap::to_external::baml_value_to_external( - &parsed, - sap.db(), - )) -} - -/// Parses the text received so far as a partial value, or `None` for no yield. -/// -/// Text that doesn't parse as the target type yet is not an error here: a reply -/// that opens with a ```` ```json ```` fence or with prose before the JSON can't -/// be coerced until the JSON starts. As in BAML v0, the stream skips that -/// partial and waits for more text; [`execute_sap_parse_final`] still reports a -/// reply that never parses. -pub fn execute_sap_parse_partial( - json: &str, - sap: &SapParseCache, - _ctx: &::sys_types::SysOpContext, -) -> Result, LlmOpError> { - // An unresolvable target type is a schema bug, not an incomplete reply. - let parse_ctx = ::bex_sap::deserializer::coercer::ParsingContext::new(sap.db()); - let target = sap - .ty_resolved() - .map_err(|err| parse_ctx.error_type_resolution(err)) - .map_err(LlmOpError::SapError)?; - - // === Jsonish === - let jsonish_options = ::bex_sap::jsonish::ParseOptions::default(); - let Ok(jsonish) = ::bex_sap::jsonish::parse(json, jsonish_options, false) else { - return Ok(None); - }; - - // === SAP parsing (a partial is the same type, parsed from the text so far) === - let Ok(parsed) = ::bex_sap::sap_model::TyResolvedRef::coerce(&parse_ctx, target, &jsonish) - else { - return Ok(None); - }; - // === Convert back to baml === - match parsed { - Some(parsed) => { - let converted = ::bex_sap::to_external::baml_value_to_external(&parsed, sap.db()); - Ok(Some(converted)) - } - None => Ok(None), - } -} - -#[cfg(test)] -mod tests { - use super::LlmOpError; - - /// A parse failure must reach BAML as `LlmClient`, which is what - /// `ai.errors.normalize` classifies into `ai.errors.ParseFailed`. Mapping - /// it to a generic internal error instead would make a provider's bad - /// answer look like an engine defect and escape the client error taxonomy. - #[test] - fn parse_response_error_maps_to_llm_client() { - let mapped: ::sys_types::VmRustFnError = - LlmOpError::ParseResponseError("bad json".to_string()).into(); - assert_eq!( - mapped, - ::sys_types::VmRustFnError::BamlError(::sys_types::VmBamlError::LlmClient { - message: "bad json".to_string(), - }) - ); - } - - /// A shape mismatch is the CALLER's problem, so it maps to - /// `InvalidArgument` — a different `baml.errors.*` class, and catchable on - /// its own. - #[test] - fn type_error_maps_to_invalid_argument() { - let mapped: ::sys_types::VmRustFnError = LlmOpError::TypeError { - expected: "string", - actual: "int".to_string(), - } - .into(); - assert_eq!( - mapped, - ::sys_types::VmRustFnError::BamlError(::sys_types::VmBamlError::InvalidArgument { - message: "expected string, got int".to_string(), - }) - ); - } -}