Test asserts in fn params and return values, test streaming

engine/baml-lib/baml-types/src/baml_value.rs

@@ -1,5 +1,5 @@
 use std::collections::HashMap;
-use std::{collections::HashSet, fmt};
+use std::{collections::{HashSet, VecDeque}, fmt};
 
 use serde::ser::{SerializeMap, SerializeSeq};
 use serde::{de::Visitor, Deserialize, Deserializer, Serialize, Serializer};
@@ -372,6 +372,10 @@ impl<T> BamlValueWithMeta<T> {
         plain_value.r#type()
     }
 
+    pub fn iter<'a>(&'a self) -> BamlValueWithMetaIterator<'a, T> {
+        BamlValueWithMetaIterator::new(self)
+    }
+
     pub fn value(self) -> BamlValue {
         match self {
             BamlValueWithMeta::String(v, _) => BamlValue::String(v),
@@ -468,6 +472,60 @@ impl<T> BamlValueWithMeta<T> {
     }
 }
 
+/// An iterator over a BamlValue and all of its sub-values.
+/// It yields entries in depth-first order.
+pub struct BamlValueWithMetaIterator<'a,T> {
+    stack: VecDeque<&'a BamlValueWithMeta<T>>,
+}
+
+impl <'a, T> BamlValueWithMetaIterator<'a, T> {
+    fn new(root: &'a BamlValueWithMeta<T>) -> Self {
+        let mut stack = VecDeque::new();
+        stack.push_back(root);
+        BamlValueWithMetaIterator { stack }
+    }
+}
+
+impl <'a,T:'a> Iterator for BamlValueWithMetaIterator<'a,T> {
+    type Item = &'a BamlValueWithMeta<T>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if let Some(value) = self.stack.pop_back() {
+            // Get all the children and push them onto the stack.
+            match value {
+                BamlValueWithMeta::String(..) | BamlValueWithMeta::Int(..) |
+                BamlValueWithMeta::Float(..) | BamlValueWithMeta::Bool(..) |
+                BamlValueWithMeta::Media(..) | BamlValueWithMeta::Enum(..) |
+                BamlValueWithMeta::Null(..) => {}
+                BamlValueWithMeta::List(items,_) => {
+                    self.stack.extend(items);
+                }
+                BamlValueWithMeta::Map(fields,_) => {
+                    for (_,v) in fields.iter() {
+                        self.stack.push_back(v);
+                    }
+                }
+                BamlValueWithMeta::Class(_, fields, _) => {
+                    for (_,v) in fields.iter() {
+                        self.stack.push_back(v);
+                    }
+                }
+            }
+            Some(&value)
+        } else {
+            None
+        }
+    }
+}
+
+impl <'a, T:'a>  IntoIterator for &'a BamlValueWithMeta<T> {
+    type Item = &'a BamlValueWithMeta<T>;
+    type IntoIter = BamlValueWithMetaIterator<'a,T>;
+    fn into_iter(self) -> Self::IntoIter {
+        self.iter()
+    }
+}
+
 impl <T> From<&BamlValueWithMeta<T>> for BamlValue {
     fn from(baml_value: &BamlValueWithMeta<T>) -> BamlValue {
         use BamlValueWithMeta::*;

engine/baml-runtime/src/internal/llm_client/mod.rs

@@ -29,8 +29,9 @@ pub type ResponseBamlValue = BamlValueWithMeta<Vec<ResponseCheck>>;
 pub fn parsed_value_to_response(baml_value: BamlValueWithFlags) -> Result<ResponseBamlValue> {
     let baml_value_with_meta: BamlValueWithMeta<Vec<(Constraint, bool)>> = baml_value.into();
     let first_failing_assert: Option<Constraint> = baml_value_with_meta
-        .meta()
         .iter()
+        .map(|v| v.meta())
+        .flatten()
         .filter_map(|(c @ Constraint { level, .. }, succeeded)| {
             if !succeeded && level == &ConstraintLevel::Assert {
                 Some(c.clone())

integ-tests/baml_src/test-files/constraints/constraints.baml

@@ -23,11 +23,6 @@ class FooAny {
 }
 
 
-class InputWithConstraint {
-  name string @assert({{this|length > 1}}, nonempty)
-  amount int @assert({{ this < 1 }}, small)
-}
-
 function PredictAge(name: string) -> FooAny {
   client GPT35
   prompt #"
@@ -43,7 +38,7 @@ function PredictAge(name: string) -> FooAny {
 }
 
 
-function PredictAgeComplex(inp: InputWithConstraint) -> InputWithConstraint {
+function PredictAgeBare(inp: string @assert({{this|length > 1}}, big_enough)) -> int @check({{this == 10102}}, too_big) {
   client GPT35
   prompt #"
     Using your understanding of the historical popularity
@@ -55,14 +50,26 @@ function PredictAgeComplex(inp: InputWithConstraint) -> InputWithConstraint {
   "#
 }
 
-function PredictAgeBare(inp: string @assert({{this|length > 1}}, big_enough)) -> int @check({{this == 10102}}, too_big) {
+function ReturnFailingAssert(inp: int @assert({{this < 10}}, small_int)) -> int @assert({{this > 100}}, big_int) {
   client GPT35
   prompt #"
-    Using your understanding of the historical popularity
-    of names, predict the age of a person with the name
-    {{ inp.name }} in years. Also predict their genus and
-    species. It's Homo sapiens (with exactly that spelling).
+    Return the next integer after {{ inp }}.
 
     {{ctx.output_format}}
   "#
 }
+
+class TwoStoriesOneTitle {
+  title string
+  story_a string @assert( {{this|length > 1000000}}, too_long_story )
+  story_b string @assert( {{this|length > 1000000}}, too_long_story )
+}
+
+function StreamFailingAssertion(theme: string, length: int) -> TwoStoriesOneTitle {
+  client GPT35
+  prompt #"
+    Tell me two different stories along the theme of {{ theme }} with the same title.
+    Please make each about {{ length }} words long.
+    {{ctx.output_format}}
+  "#
+}

integ-tests/openapi/baml_client/openapi.yaml

@@ -610,19 +610,6 @@ paths:
                 - checks
                 additionalProperties: false
       operationId: PredictAgeBare
-  /call/PredictAgeComplex:
-    post:
-      requestBody:
-        $ref: '#/components/requestBodies/PredictAgeComplex'
-      responses:
-        '200':
-          description: Successful operation
-          content:
-            application/json:
-              schema:
-                title: PredictAgeComplexResponse
-                $ref: '#/components/schemas/InputWithConstraint'
-      operationId: PredictAgeComplex
   /call/PromptTestClaude:
     post:
       requestBody:
@@ -714,6 +701,19 @@ paths:
                 title: PromptTestStreamingResponse
                 type: string
       operationId: PromptTestStreaming
+  /call/ReturnFailingAssert:
+    post:
+      requestBody:
+        $ref: '#/components/requestBodies/ReturnFailingAssert'
+      responses:
+        '200':
+          description: Successful operation
+          content:
+            application/json:
+              schema:
+                title: ReturnFailingAssertResponse
+                type: integer
+      operationId: ReturnFailingAssert
   /call/SchemaDescriptions:
     post:
       requestBody:
@@ -740,6 +740,19 @@ paths:
                 title: StreamBigNumbersResponse
                 $ref: '#/components/schemas/BigNumbers'
       operationId: StreamBigNumbers
+  /call/StreamFailingAssertion:
+    post:
+      requestBody:
+        $ref: '#/components/requestBodies/StreamFailingAssertion'
+      responses:
+        '200':
+          description: Successful operation
+          content:
+            application/json:
+              schema:
+                title: StreamFailingAssertionResponse
+                $ref: '#/components/schemas/TwoStoriesOneTitle'
+      operationId: StreamFailingAssertion
   /call/StreamOneBigNumber:
     post:
       requestBody:
@@ -1885,22 +1898,6 @@ components:
             required:
             - inp
             additionalProperties: false
-    PredictAgeComplex:
-      required: true
-      content:
-        application/json:
-          schema:
-            title: PredictAgeComplexRequest
-            type: object
-            properties:
-              inp:
-                $ref: '#/components/schemas/InputWithConstraint'
-              __baml_options__:
-                nullable: true
-                $ref: '#/components/schemas/BamlOptions'
-            required:
-            - inp
-            additionalProperties: false
     PromptTestClaude:
       required: true
       content:
@@ -2013,6 +2010,22 @@ components:
             required:
             - input
             additionalProperties: false
+    ReturnFailingAssert:
+      required: true
+      content:
+        application/json:
+          schema:
+            title: ReturnFailingAssertRequest
+            type: object
+            properties:
+              inp:
+                type: integer
+              __baml_options__:
+                nullable: true
+                $ref: '#/components/schemas/BamlOptions'
+            required:
+            - inp
+            additionalProperties: false
     SchemaDescriptions:
       required: true
       content:
@@ -2045,6 +2058,25 @@ components:
             required:
             - digits
             additionalProperties: false
+    StreamFailingAssertion:
+      required: true
+      content:
+        application/json:
+          schema:
+            title: StreamFailingAssertionRequest
+            type: object
+            properties:
+              theme:
+                type: string
+              length:
+                type: integer
+              __baml_options__:
+                nullable: true
+                $ref: '#/components/schemas/BamlOptions'
+            required:
+            - theme
+            - length
+            additionalProperties: false
     StreamOneBigNumber:
       required: true
       content:
@@ -2941,13 +2973,13 @@ components:
             checks:
               type: object
               properties:
-                no_infants:
-                  $ref: '#components/schemas/Check'
                 earth_aged:
                   $ref: '#components/schemas/Check'
+                no_infants:
+                  $ref: '#components/schemas/Check'
               required:
-              - no_infants
               - earth_aged
+              - no_infants
               additionalProperties: false
           required:
           - value
@@ -3091,15 +3123,15 @@ components:
             checks:
               type: object
               properties:
-                trivial:
-                  $ref: '#components/schemas/Check'
                 regex_bad:
                   $ref: '#components/schemas/Check'
+                trivial:
+                  $ref: '#components/schemas/Check'
                 regex_good:
                   $ref: '#components/schemas/Check'
               required:
-              - trivial
               - regex_bad
+              - trivial
               - regex_good
               additionalProperties: false
           required:
@@ -3158,17 +3190,6 @@ components:
       - prop2
       - prop3
       additionalProperties: false
-    InputWithConstraint:
-      type: object
-      properties:
-        name:
-          type: string
-        amount:
-          type: integer
-      required:
-      - name
-      - amount
-      additionalProperties: false
     Martian:
       type: object
       properties:
@@ -3542,6 +3563,20 @@ components:
       - prop1
       - prop2
       additionalProperties: false
+    TwoStoriesOneTitle:
+      type: object
+      properties:
+        title:
+          type: string
+        story_a:
+          type: string
+        story_b:
+          type: string
+      required:
+      - title
+      - story_a
+      - story_b
+      additionalProperties: false
     UnionTest_ReturnType:
       type: object
       properties: