Addressed issue where dynamic types were not streaming. (#1606)

- Updated `ir_helpers/mod.rs` to introduce `IRHelperExtended` and
`IRSemanticStreamingHelper` traits, adding methods for class field
management and streaming behavior.
- Modified `mod.rs` to export the new traits and updated relevant
imports across the codebase.
- Refactored `validate_streaming_state` and `process_node` functions in
`semantic_streaming.rs` to utilize the new traits for improved type
handling.
- Adjusted various modules to implement the extended functionality,
ensuring better integration with runtime context and type management.

This update improves the flexibility and usability of the IR helper
functions, facilitating better handling of class fields and streaming
requirements.

<!-- ELLIPSIS_HIDDEN -->


----

> [!IMPORTANT]
> Introduce `IRHelperExtended` and `IRSemanticStreamingHelper` traits to
enhance dynamic type handling and streaming behavior in the Intermediate
Representation.
> 
>   - **Traits**:
> - Introduce `IRHelperExtended` and `IRSemanticStreamingHelper` in
`ir_helpers/mod.rs` for class field management and streaming behavior.
>     - Implement these traits in `IntermediateRepr` and `ScopedIr`.
>   - **Functions**:
> - Refactor `validate_streaming_state` and `process_node` in
`semantic_streaming.rs` to use new traits.
> - Update `parsed_value_to_response` in `llm_client/mod.rs` and
`helpers/mod.rs` to utilize new traits.
>   - **Modules**:
> - Modify `mod.rs` to export new traits and update imports across the
codebase.
> - Adjust various modules to implement extended functionality for
better runtime context and type management.
> 
> <sup>This description was created by </sup>[<img alt="Ellipsis"
src="https://img.shields.io/badge/Ellipsis-blue?color=175173">](https://www.ellipsis.dev?ref=BoundaryML%2Fbaml&utm_source=github&utm_medium=referral)<sup>
for e9cb26c. It will automatically
update as commits are pushed.</sup>


<!-- ELLIPSIS_HIDDEN -->

Author: hellovai

engine/baml-lib/baml-core/src/ir/ir_helpers/mod.rs

@@ -7,7 +7,7 @@ use std::path::PathBuf;
 
 use crate::ir::IntermediateRepr;
 
-use super::{scope_diagnostics::ScopeStack, IRHelper};
+use super::{scope_diagnostics::ScopeStack, IRHelper, IRHelperExtended};
 use crate::ir::jinja_helpers::evaluate_predicate;
 
 #[derive(Default)]

engine/baml-lib/baml-core/src/ir/ir_helpers/to_baml_arg.rs

@@ -1,12 +1,13 @@
-mod ir_helpers;
+pub mod ir_helpers;
 pub mod jinja_helpers;
 mod json_schema;
 pub mod repr;
 mod walker;
 
 pub use ir_helpers::{
     scope_diagnostics, ArgCoercer, ClassFieldWalker, ClassWalker, ClientWalker, EnumValueWalker,
-    EnumWalker, FunctionWalker, IRHelper, RetryPolicyWalker, TemplateStringWalker, TestCaseWalker,
+    EnumWalker, FunctionWalker, IRHelper, IRHelperExtended, IRSemanticStreamingHelper,
+    RetryPolicyWalker, TemplateStringWalker, TestCaseWalker, TypeAliasWalker,
 };
 
 pub(super) use repr::IntermediateRepr;

engine/baml-lib/baml-core/src/ir/mod.rs

@@ -5,7 +5,7 @@ use crate::deserializer::coercer::ParsingError;
 use crate::{BamlValueWithFlags, Flag};
 use indexmap::{IndexMap, IndexSet};
 use internal_baml_core::ir::repr::{IntermediateRepr, Walker};
-use internal_baml_core::ir::{Field, IRHelper};
+use internal_baml_core::ir::{Field, IRHelper, IRHelperExtended, IRSemanticStreamingHelper};
 
 use baml_types::{
     BamlMap, BamlValueWithMeta, Completion, CompletionState, FieldType, ResponseCheck,
@@ -31,7 +31,7 @@ pub enum StreamingError {
 /// For a given baml value, traverse its nodes, comparing the completion state
 /// of each node against the streaming behavior of the node's type.
 pub fn validate_streaming_state(
-    ir: &IntermediateRepr,
+    ir: &impl IRHelperExtended,
     baml_value: &BamlValueWithFlags,
     field_type: &FieldType,
     allow_partials: bool,
@@ -63,7 +63,7 @@ pub fn validate_streaming_state(
 ///   allow_partials: Whether this node may contain partial values. (Once we
 ///                   see a false, all child nodes will also get false).
 fn process_node(
-    ir: &IntermediateRepr,
+    ir: &impl IRHelperExtended,
     value: BamlValueWithMeta<(CompletionState, &FieldType)>,
     allow_partials: bool,
     depth: usize,
@@ -233,15 +233,9 @@ fn process_node(
 
 /// Extract the field names from a field_type that is expected to be a `Class`.
 /// If it is not a known class, return no field names.
-fn type_field_names(ir: &IntermediateRepr, field_type: &FieldType) -> IndexSet<String> {
+fn type_field_names(ir: &impl IRHelperExtended, field_type: &FieldType) -> IndexSet<String> {
     match ir.distribute_metadata(field_type).0 {
-        FieldType::Class(class_name) => match ir.find_class(class_name) {
-            Err(_) => IndexSet::new(),
-            Ok(class) => class
-                .walk_fields()
-                .map(|field| field.name().to_string())
-                .collect(),
-        },
+        FieldType::Class(class_name) => ir.class_field_names(class_name).unwrap_or_default(),
         _ => IndexSet::new(),
     }
 }
@@ -250,31 +244,15 @@ fn type_field_names(ir: &IntermediateRepr, field_type: &FieldType) -> IndexSet<S
 /// fields in the class need to be filled in by a null. A field needs to be
 /// filled by a null if it is not present in the map value.
 fn fields_needing_null_filler<'a>(
-    ir: &'a IntermediateRepr,
+    ir: &'a impl IRSemanticStreamingHelper,
     class_name: &'a str,
     value_names: HashSet<String>,
     allow_partials: bool,
 ) -> Result<HashSet<String>, anyhow::Error> {
     if allow_partials == false {
         return Ok(HashSet::new());
     }
-    let res = match ir.find_class(class_name) {
-        Err(_) => Ok(HashSet::new()),
-        Ok(class) => {
-            let missing_fields = class
-                .walk_fields()
-                .filter_map(|field: Walker<'_, &Field>| {
-                    if !value_names.contains(field.name()) {
-                        Some(field.name().to_string())
-                    } else {
-                        None
-                    }
-                })
-                .collect();
-            Ok(missing_fields)
-        }
-    };
-    res
+    ir.find_class_fields_needing_null_filler(class_name, &value_names)
 }
 
 /// For a given type, assume that it is a class, and list the fields of that
@@ -285,33 +263,21 @@ fn fields_needing_null_filler<'a>(
 /// and return an empty set (because we are ignoring the "@stream.not_null" property,
 /// which only applies when `allow_partials==true`).
 fn needed_fields(
-    ir: &IntermediateRepr,
+    ir: &impl IRHelperExtended,
     class_name: &str,
     allow_partials: bool,
 ) -> Result<HashSet<String>, anyhow::Error> {
     if allow_partials == false {
         return Ok(HashSet::new());
     }
-    let class = ir
-        .find_class(class_name)
+    ir.class_streaming_needed_fields(class_name)
         .map_err(|_| StreamingError::ExpectedClass)
-        .context("needed_fields failed to lookup class")?;
-    let needed_fields = class
-        .walk_fields()
-        .filter_map(|field: Walker<'_, &Field>| {
-            if field.streaming_needed() {
-                Some(field.name().to_string())
-            } else {
-                None
-            }
-        })
-        .collect();
-    Ok(needed_fields)
+        .context("needed_fields failed to lookup class")
 }
 
 /// Whether a type must be complete before being included as a node
 /// in a streamed value.
-fn required_done(ir: &IntermediateRepr, field_type: &FieldType) -> bool {
+fn required_done(ir: &impl IRHelperExtended, field_type: &FieldType) -> bool {
     let (base_type, (_, streaming_behavior)) = ir.distribute_metadata(field_type);
     let type_implies_done = match base_type {
         FieldType::Primitive(tv) => match tv {
@@ -353,7 +319,7 @@ fn completion_state(flags: &Vec<Flag>) -> CompletionState {
     }
 }
 
-fn type_streaming_behavior(ir: &IntermediateRepr, r#type: &FieldType) -> StreamingBehavior {
+fn type_streaming_behavior(ir: &impl IRHelperExtended, r#type: &FieldType) -> StreamingBehavior {
     let (_base_type, (_constraints, streaming_behavior)) = ir.distribute_metadata(r#type);
     streaming_behavior
 }

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

@@ -2,9 +2,10 @@ pub mod common;
 use std::{collections::HashSet, path::PathBuf};
 
 use anyhow::Result;
-use baml_types::{EvaluationContext, JinjaExpression};
 use baml_types::{BamlValueWithMeta, ResponseCheck, StreamingBehavior};
+use baml_types::{EvaluationContext, JinjaExpression};
 use indexmap::{IndexMap, IndexSet};
+use internal_baml_core::ir::IRHelperExtended;
 use internal_baml_core::{
     ast::Field,
     internal_baml_diagnostics::SourceFile,
@@ -263,7 +264,6 @@ pub fn parsed_value_to_response(
     field_type: &FieldType,
     allow_partials: bool,
 ) -> Result<ResponseBamlValue> {
-
     let meta_flags: BamlValueWithMeta<Vec<Flag>> = baml_value.clone().into();
     let baml_value_with_meta: BamlValueWithMeta<Vec<(String, JinjaExpression, bool)>> =
         baml_value.clone().into();
@@ -292,6 +292,6 @@ pub fn parsed_value_to_response(
     let response_value = baml_value_with_streaming
         .zip_meta(&value_with_response_checks)?
         .zip_meta(&meta_flags)?
-        .map_meta(|((x, y), z)| (z.clone(), y.clone(), x.clone() ));
+        .map_meta(|((x, y), z)| (z.clone(), y.clone(), x.clone()));
     Ok(ResponseBamlValue(response_value))
 }

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

@@ -12,7 +12,7 @@ pub mod traits;
 use anyhow::{Context, Result};
 
 use baml_types::{BamlMap, BamlValueWithMeta, FieldType, JinjaExpression, ResponseCheck};
-use internal_baml_core::ir::{repr::IntermediateRepr, ClientWalker};
+use internal_baml_core::ir::{repr::IntermediateRepr, ClientWalker, IRHelper, IRHelperExtended};
 use internal_baml_jinja::RenderedPrompt;
 use internal_llm_client::AllowedRoleMetadata;
 pub use jsonish::ResponseBamlValue;
@@ -33,7 +33,7 @@ use wasm_bindgen::JsValue;
 
 /// Validate a parsed value, checking asserts and checks.
 pub fn parsed_value_to_response(
-    ir: &IntermediateRepr,
+    ir: &impl IRHelperExtended,
     baml_value: BamlValueWithFlags,
     field_type: &FieldType,
     allow_partials: bool,

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

@@ -1,20 +1,25 @@
 mod render_output_format;
+pub(crate) mod scoped_ir;
 use internal_llm_client::ClientSpec;
 use jsonish::{BamlValueWithFlags, ResponseBamlValue};
 use render_output_format::render_output_format;
 
 use anyhow::Result;
-use baml_types::{BamlValue, FieldType};
+use baml_types::{BamlValue, FieldType, StreamingBehavior};
 use internal_baml_core::{
     error_unsupported,
-    ir::{repr::IntermediateRepr, FunctionWalker, IRHelper},
+    ir::{
+        repr::IntermediateRepr, FunctionWalker, IRHelper, IRHelperExtended,
+        IRSemanticStreamingHelper,
+    },
 };
 use internal_baml_jinja::{
     types::OutputFormatContent, RenderContext, RenderContext_Client, RenderedPrompt,
     TemplateStringMacro,
 };
+use scoped_ir::ScopedIr;
 
-use crate::RuntimeContext;
+use crate::{runtime_context::RuntimeClassOverride, RuntimeContext};
 
 use super::llm_client::parsed_value_to_response;
 
@@ -55,6 +60,7 @@ impl PromptRenderer {
     pub fn parse(
         &self,
         ir: &IntermediateRepr,
+        ctx: &RuntimeContext,
         raw_string: &str,
         allow_partials: bool,
     ) -> Result<ResponseBamlValue> {
@@ -64,7 +70,8 @@ impl PromptRenderer {
             raw_string,
             allow_partials,
         )?;
-        let res = parsed_value_to_response(ir, parsed, &self.output_type, allow_partials);
+        let scoped_ir = ScopedIr::new(ir, ctx);
+        let res = parsed_value_to_response(&scoped_ir, parsed, &self.output_type, allow_partials);
         res
     }
 

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

@@ -4,7 +4,7 @@ use anyhow::Result;
 use baml_types::BamlValue;
 use indexmap::{IndexMap, IndexSet};
 use internal_baml_core::ir::{
-    repr::IntermediateRepr, ClassWalker, EnumWalker, FieldType, IRHelper,
+    repr::IntermediateRepr, ClassWalker, EnumWalker, FieldType, IRHelper, IRHelperExtended,
 };
 use internal_baml_jinja::types::{Class, Enum, Name, OutputFormatContent};