add test for streaming failure for lists (#2092)

Co-authored-by: hellovai <vbv@boundaryml.com>

Author: aaronvg

.claude/settings.local.json

@@ -0,0 +1,15 @@
+{
+  "permissions": {
+    "allow": [
+      "Bash(grep:*)",
+      "Bash(uv run maturin develop:*)",
+      "Bash(BAML_LOG_INTERNAL=debug BAML_INTERNAL_LOG=debug infisical run --env=test -- uv run pytest tests/test_functions.py::test_semantic_streaming -s)",
+      "Bash(python test:*)",
+      "Bash(uv run:*)",
+      "Bash(find:*)",
+      "Bash(rg:*)",
+      "Bash(cargo check:*)"
+    ],
+    "deny": []
+  }
+}

CONTRIBUTING.md

@@ -77,6 +77,8 @@ Issues labeled `good first issue` are a good place to start.
 
 ### Steps to Build and Test Locally
 
+1. Install protobuf generator for Go `brew install protoc-gen-go`
+
 1. Install Rust
 
 2. Run `cargo build` in `engine/` and make sure everything builds on your machine.

engine/baml-lib/jsonish/src/deserializer/coercer/field_type.rs

@@ -124,7 +124,18 @@ impl TypeCoercer for TypeIR {
             }
         }
         if let Some(CompletionState::Incomplete) = value.map(|v| v.completion_state()) {
-            result.iter_mut().for_each(|v| v.add_flag(Flag::Incomplete));
+            match result {
+                Ok(mut v) => {
+                    if self.meta().streaming_behavior.done
+                        && ctx.do_not_use_mode == baml_types::StreamingMode::Streaming
+                    {
+                        return Err(ctx.error_internal("Streaming field is not done"));
+                    }
+                    v.add_flag(Flag::Incomplete);
+                    return Ok(v);
+                }
+                Err(e) => return Err(e),
+            }
         }
         result
     }

engine/baml-lib/jsonish/src/deserializer/coercer/mod.rs

@@ -21,6 +21,8 @@ use crate::jsonish;
 pub struct ParsingContext<'a> {
     pub scope: Vec<String>,
     visited: HashSet<(String, jsonish::Value)>,
+    /// THIS IS A TEMPORARY HACK (ask vaibhav)
+    pub do_not_use_mode: baml_types::StreamingMode,
     pub of: &'a OutputFormatContent,
 }
 
@@ -32,10 +34,14 @@ impl ParsingContext<'_> {
         self.scope.join(".")
     }
 
-    pub(crate) fn new(of: &OutputFormatContent) -> ParsingContext<'_> {
+    pub(crate) fn new(
+        of: &OutputFormatContent,
+        mode: baml_types::StreamingMode,
+    ) -> ParsingContext<'_> {
         ParsingContext {
             scope: Vec::new(),
             visited: HashSet::new(),
+            do_not_use_mode: mode,
             of,
         }
     }
@@ -47,6 +53,7 @@ impl ParsingContext<'_> {
             scope: new_scope,
             visited: self.visited.clone(),
             of: self.of,
+            do_not_use_mode: self.do_not_use_mode.clone(),
         }
     }
 
@@ -63,6 +70,7 @@ impl ParsingContext<'_> {
             scope: self.scope.clone(),
             visited: new_visited,
             of: self.of,
+            do_not_use_mode: self.do_not_use_mode.clone(),
         }
     }
 

engine/baml-lib/jsonish/src/helpers/mod.rs

@@ -237,7 +237,7 @@ fn relevant_data_models<'a>(
                         .map(|field| find_existing_class_field(name, &field, &walker, env_values))
                         .map(|field| {
                             let (name, t, prop1, prop2) = field?;
-                            let t = if partialize {
+                            let t = if partialize && !metadata.streaming_behavior.done {
                                 t.to_streaming_type(ir).to_ir_type()
                             } else {
                                 t

engine/baml-lib/jsonish/src/lib.rs

@@ -245,7 +245,14 @@ pub fn from_str(
 
     // Pick the schema that is the most specific.
     log::debug!("Parsed JSONish (step 1 of parsing) {raw_string_is_done}: {value:#?}");
-    let ctx = ParsingContext::new(of);
+    let ctx = ParsingContext::new(
+        of,
+        if raw_string_is_done {
+            baml_types::StreamingMode::NonStreaming
+        } else {
+            baml_types::StreamingMode::Streaming
+        },
+    );
 
     // Determine the best way to get the desired schema from the parsed schema.
 

engine/baml-lib/jsonish/src/tests/macros.rs

@@ -201,19 +201,19 @@ macro_rules! test_partial_deserializer_streaming_failure {
 
             let parsed = from_str(&target, &target_type, $raw_string, false);
 
-            assert!(parsed.is_ok(), "Failed to parse: {:?}", parsed);
-
-            let result = crate::helpers::parsed_value_to_response(
-                &ir,
-                parsed.unwrap(),
-                baml_types::StreamingMode::Streaming,
-            );
-
-            assert!(
-                result.is_err(),
-                "Failed not to parse: {}",
-                json!(result.unwrap().serialize_partial())
-            );
+            assert!(parsed.is_err(), "Got a result: {:?}", parsed);
+
+            // let result = crate::helpers::parsed_value_to_response(
+            //     &ir,
+            //     parsed.unwrap(),
+            //     baml_types::StreamingMode::Streaming,
+            // );
+
+            // assert!(
+            //     result.is_err(),
+            //     "Failed not to parse: {}",
+            //     json!(result.unwrap().serialize_partial())
+            // );
         }
     };
 }

engine/baml-lib/jsonish/src/tests/test_class_2.rs

@@ -527,3 +527,232 @@ Actions:
     }
   ]
 );
+
+const TEST_FILE_STREAMING_NOT_NULL_LIST: &str = r###"
+class SemanticContainer2 {    
+    three_small_things SmallThing[] @description("Should have three items.")
+}
+
+class SemanticContainer {
+    sixteen_digit_number int
+    string_with_twenty_words string @stream.done
+    class_1 ClassWithoutDone
+    class_2 ClassWithBlockDone
+    class_done_needed ClassWithBlockDone @stream.not_null
+    class_needed ClassWithoutDone @stream.not_null
+    three_small_things SmallThing[] @description("Should have three items.")
+    final_string string
+}
+
+class ClassWithoutDone {
+    i_16_digits int
+    s_20_words string @description("A string with 20 words in it") @stream.with_state
+}
+
+class ClassWithBlockDone {
+    i_16_digits int
+    s_20_words string
+    @@stream.done
+}
+
+class SmallThing {
+    i_16_digits int @stream.not_null
+    i_8_digits int
+}
+
+function MakeSemanticContainer() -> SemanticContainer {
+    client "openai/gpt-4o"
+    prompt #"
+        {{ ctx.output_format }}
+    "#
+}
+"###;
+
+test_deserializer!(
+  test_streaming_not_null_list,
+  TEST_FILE_STREAMING_NOT_NULL_LIST,
+  r#"
+      {
+      "sixteen_digit_number": 1234567890123456,
+      "string_with_twenty_words": "This is an example string that contains exactly twenty words for the requested output in JSON format. It is quite informative.",
+      "class_1": {
+        "i_16_digits": 9876543210123456,
+        "s_20_words": "Another example string used for class one with twenty distinct words to complete the requirements of this JSON structure."
+      },
+      "class_2": {
+        "i_16_digits": 4567890123456789,
+        "s_20_words": "Class two also has a longer description that encompasses twenty unique words to fulfill the output specifications."
+      },
+      "class_done_needed": {
+        "i_16_digits": 7890123456789012,
+        "s_20_words": "This class indicates what has been completed and also includes twenty words in its description for clarity."
+      },
+      "class_needed": {
+        "i_16_digits": 3210987654321098,
+        "s_20_words": "Another class that explains what is currently needed for completion includes another twenty-word explanation for better understanding."
+      },
+      "three_small_things": [
+        {
+          "i_16_digits": 1357924680135792,
+          "i_8_digits": 24681357
+        },
+        {
+          "i_16_digits": 2468135790246813,
+          "i_8_digits": 86421357
+        },
+        {
+          "i_16_digits": 3571598642035791,
+          "i_8_digits": 97586421
+        }
+      ],
+      "final_string": "This final string provides a brief conclusion or summary of the entire JSON formatted data provided above."
+  }  
+  "#,
+  TypeIR::class("SemanticContainer"),
+  {
+    "sixteen_digit_number": 1234567890123456i64,
+    "string_with_twenty_words": "This is an example string that contains exactly twenty words for the requested output in JSON format. It is quite informative.",
+    "class_1": {
+      "i_16_digits": 9876543210123456i64,
+      "s_20_words": "Another example string used for class one with twenty distinct words to complete the requirements of this JSON structure."
+    },
+    "class_2": {
+      "i_16_digits": 4567890123456789i64,
+      "s_20_words": "Class two also has a longer description that encompasses twenty unique words to fulfill the output specifications."
+    },
+    "class_done_needed": {
+      "i_16_digits": 7890123456789012i64,
+      "s_20_words": "This class indicates what has been completed and also includes twenty words in its description for clarity."
+    },
+    "class_needed": {
+      "i_16_digits": 3210987654321098i64,
+      "s_20_words": "Another class that explains what is currently needed for completion includes another twenty-word explanation for better understanding."
+    },
+    "three_small_things": [
+      {
+        "i_16_digits": 1357924680135792i64,
+        "i_8_digits": 24681357
+      },
+      {
+        "i_16_digits": 2468135790246813i64,
+        "i_8_digits": 86421357
+      },
+      {
+        "i_16_digits": 3571598642035791i64,
+        "i_8_digits": 97586421
+      }
+    ],
+    "final_string": "This final string provides a brief conclusion or summary of the entire JSON formatted data provided above."
+  }
+);
+
+test_partial_deserializer!(
+  test_streaming_not_null_list_partial,
+  TEST_FILE_STREAMING_NOT_NULL_LIST,
+  r#"
+      {
+      "sixteen_digit_number": 1234567890123456,
+      "string_with_twenty_words": "This is an example string that contains exactly twenty words for the requested output in JSON format. It is quite informative.",
+      "class_1": {
+        "i_16_digits": 9876543210123456,
+        "s_20_words": "Another example string used for class one with twenty distinct words to complete the requirements of this JSON structure."
+      },
+      "class_2": {
+        "i_16_digits": 4567890123456789,
+        "s_20_words": "Class two also has a longer description that encompasses twenty unique words to fulfill the output specifications."
+      },
+      "class_done_needed": {
+        "i_16_digits": 7890123456789012,
+        "s_20_words": "This class indicates what has been completed and also includes twenty words in its description for clarity."
+      },
+      "class_needed": {
+        "i_16_digits": 3210987654321098,
+        "s_20_words": "Another class that explains what is currently needed for completion includes another twenty-word explanation for better understanding."
+      },
+      "three_small_things": [{
+        "i_16_digits": 123
+        
+  "#,
+  TypeIR::class("SemanticContainer"),
+  {
+    "sixteen_digit_number": 1234567890123456i64,
+    "string_with_twenty_words": "This is an example string that contains exactly twenty words for the requested output in JSON format. It is quite informative.",
+    "class_1": {
+      "i_16_digits": 9876543210123456i64,
+      "s_20_words": "Another example string used for class one with twenty distinct words to complete the requirements of this JSON structure."
+    },
+    "class_2": {
+      "i_16_digits": 4567890123456789i64,
+      "s_20_words": "Class two also has a longer description that encompasses twenty unique words to fulfill the output specifications."
+    },
+    "class_done_needed": {
+      "i_16_digits": 7890123456789012i64,
+      "s_20_words": "This class indicates what has been completed and also includes twenty words in its description for clarity."
+    },
+    "class_needed": {
+      "i_16_digits": 3210987654321098i64,
+      "s_20_words": "Another class that explains what is currently needed for completion includes another twenty-word explanation for better understanding."
+    },
+    "three_small_things": [
+    ],
+    "final_string": null
+  }
+);
+
+test_deserializer!(
+  test_streaming_not_null_list_partial_2,
+  TEST_FILE_STREAMING_NOT_NULL_LIST,
+  r#"
+      {
+        "i_16_digits": 123456789012345,
+        "i_8_digits": 12345
+      }
+  "#,
+  TypeIR::class("SmallThing"),
+  {
+    "i_16_digits": 123456789012345i64,
+    "i_8_digits": 12345
+  }
+);
+
+test_partial_deserializer!(
+    test_streaming_not_null_list_partial_3_0,
+    TEST_FILE_STREAMING_NOT_NULL_LIST,
+    r#"
+      {
+        "three_small_things": [
+          {
+            "i_16_digits": 123456789012345
+  "#,
+  TypeIR::class("SemanticContainer2"),
+  // we should return an empty list here, not null
+  // SmallThing.i_16_digits is not_null, so the entry is not valid
+  // until the number is complete
+  {
+    "three_small_things": [
+    ]
+  }
+);
+
+test_partial_deserializer!(
+  test_streaming_not_null_list_partial_3_1,
+  TEST_FILE_STREAMING_NOT_NULL_LIST,
+  r#"
+    {
+      "three_small_things": [
+        {
+          "i_16_digits": 123456789012345,
+"#,
+TypeIR::class("SemanticContainer2"),
+// we should return an empty list here, not null
+// SmallThing.i_16_digits is not_null, so the entry is not valid
+// until the number is complete
+{
+  "three_small_things": [
+    {
+      "i_16_digits": 123456789012345i64,
+      "i_8_digits": null
+    }
+  ]
+}
+);