Fix regression in ts streaming codegen (#2533)

Author: imalsogreg

engine/baml-lib/baml-types/src/ir_type/mod.rs

@@ -1647,4 +1647,38 @@ mod tests {
         );
         assert_eq!(alias.to_streaming_type(&TestLookup), expected);
     }
+
+    #[test]
+    fn partialize_mixed_done_union() {
+        let mut done_variant = TypeIR::class("FooDone");
+        done_variant.meta_mut().streaming_behavior.done = true;
+
+        let streamable_variant = TypeIR::class("MessageToUser");
+        let union = TypeIR::Union(
+            unsafe { UnionTypeGeneric::new_unsafe(vec![done_variant, streamable_variant]) },
+            Default::default(),
+        );
+        let streaming_type = union.to_streaming_type(&TestLookup);
+        let streaming_type_variants: Vec<TypeStreaming> = match streaming_type {
+            TypeStreaming::Union(union, _) => union.view().flatten(),
+            _ => panic!("Expected union"),
+        };
+        assert_eq!(streaming_type_variants.len(), 3);
+
+        let mut expected_first_variant = TypeStreaming::class("FooDone");
+        expected_first_variant.meta_mut().streaming_behavior.done = true;
+
+        let expected_second_variant = TypeStreaming::Class {
+            name: "MessageToUser".to_string(),
+            mode: StreamingMode::Streaming,
+            dynamic: false,
+            meta: Default::default(),
+        };
+
+        dbg!(&streaming_type_variants[0]);
+        dbg!(&streaming_type_variants[1]);
+        dbg!(&streaming_type_variants[2]);
+        assert_eq!(streaming_type_variants[0], expected_first_variant);
+        assert_eq!(streaming_type_variants[1], expected_second_variant);
+    }
 }

engine/generators/languages/typescript/src/type.rs

@@ -260,7 +260,16 @@ impl SerializeType for TypeTS {
             TypeTS::TypeAlias { package, name, .. } => {
                 format!("{}{}", package.relative_from(pkg), name)
             }
-            TypeTS::Union { .. } => self.default_name_within_union(),
+            TypeTS::Union { variants, .. } => {
+                let mut parts: Vec<String> = Vec::with_capacity(variants.len());
+                for variant in variants.iter() {
+                    let rendered = variant.serialize_type(pkg);
+                    if !parts.contains(&rendered) {
+                        parts.push(rendered);
+                    }
+                }
+                parts.join(" | ")
+            }
             TypeTS::Enum {
                 package,
                 name,
@@ -274,7 +283,7 @@ impl SerializeType for TypeTS {
                 }
             }
             TypeTS::List(inner, _) => match &**inner {
-                TypeTS::Union { .. } => format!("({})[]", inner.default_name_within_union()),
+                TypeTS::Union { .. } => format!("({})[]", inner.serialize_type(pkg)),
                 _ => {
                     if inner.meta().is_optional() {
                         format!("({})[]", inner.serialize_type(pkg))

flake.nix

@@ -46,8 +46,8 @@
         buildInputs = (with pkgs; [
           cmake
           git
-        go
-        gotools
+          go
+          gotools
           mise
           openssl
           pkg-config

integ-tests/baml_src/test-files/tools/todo-llm.baml

@@ -0,0 +1,33 @@
+class AddTodoItem {
+  type "add_todo_item" @stream.not_null
+  item string
+  time string
+  description string @description("20 word description of the item")
+  @@stream.done
+}
+
+class TodoMessageToUser {
+  type "todo_message_to_user" @stream.not_null
+  message string @description("A message to the user, about 50 words long")
+}
+
+type TodoTool = AddTodoItem | TodoMessageToUser
+
+function ChooseTodoTools(query: string) -> TodoTool[] {
+  client GPT4
+  prompt #"
+    Choose tools to satisfy the user query.
+    For example, if they ask for "5 todo items for learning chess",
+    return a list of 5 "add_todo_item" objects and single "todo_message_to_user"
+    object. All requests should end with a "todo_message_to_user" object.
+
+    {{ ctx.output_format }}
+  "#
+}
+
+test ChooseTodoTools {
+  functions [ChooseTodoTools]
+  args {
+    query "5 todo items for learning chess"
+  }
+}