Skip to content

Commit ff2e76e

Browse files
imalsogregimalsogreg
authored
Superficial streaming for workflows (#1832)
When a workflow ends with an LLM call, the workflow will stream that LLM call. In other words, when you call a function whose final expression is a single LLM call, the result can be streamed. But if the final expression is something other than an LLM call, like a list or a class constructor, the result will be delivered in one shot. <!-- ELLIPSIS_HIDDEN --> ---- > [!IMPORTANT] > This PR enables streaming for workflows ending with an LLM call by evaluating if the final expression is a single LLM call, with updates to type inference, expression evaluation, and runtime execution. > > - **Behavior**: > - Workflows ending with an LLM call now support streaming if the final expression is a single LLM call. > - Non-LLM final expressions (e.g., lists, class constructors) are delivered in one shot. > - **Type Inference**: > - Updates to `infer_types_in_context()` in `expr_typecheck.rs` to handle new expression types. > - Introduces `initial_typing_context()` in `repr.rs` for better type inference. > - **Expression Evaluation**: > - Adds `eval_to_value_or_llm_call()` in `eval_expr.rs` to evaluate expressions to either a value or an LLM call. > - Modifies `beta_reduce()` in `eval_expr.rs` to support LLM function evaluation. > - **Runtime Execution**: > - Updates `BamlRuntime` in `lib.rs` to handle expression functions and determine if they result in an LLM call. > - Introduces `expr_eval_result()` in `lib.rs` to evaluate expression functions and determine streaming capability. > > <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 002d33f. You can [customize](https://app.ellipsis.dev/BoundaryML/settings/summaries) this summary. It will automatically update as commits are pushed.</sup> <!-- ELLIPSIS_HIDDEN -->
1 parent 960c7d8 commit ff2e76e

File tree

4 files changed

+668
-248
lines changed

4 files changed

+668
-248
lines changed

‎engine/baml-lib/baml-core/src/ir/repr.rs

+9-10
Original file line numberDiff line numberDiff line change
@@ -171,10 +171,6 @@ impl WithRepr<ExprFunction> for ExprFnWalker<'_> {
171171
}
172172
}
173173

174-
fn weird_default() -> FieldType {
175-
FieldType::Primitive(TypeValue::Null)
176-
}
177-
178174
impl WithRepr<Function> for ExprFnWalker<'_> {
179175
fn repr(&self, db: &ParserDatabase) -> Result<Function> {
180176
// TODO: Drop weird default (replace by better validation).
@@ -626,10 +622,10 @@ impl IntermediateRepr {
626622
repr.retry_policies
627623
.sort_by(|a, b| a.elem.name.0.cmp(&b.elem.name.0));
628624

629-
// TODO: Necessary?
625+
let mut typing_context = initial_typing_context(&repr);
630626
for expr_fn in repr.expr_fns.iter_mut() {
631627
let expr = expr_fn.elem.expr.clone();
632-
let inferred_expr = infer_types_in_context(&mut HashMap::new(), Arc::new(expr));
628+
let inferred_expr = infer_types_in_context(&mut typing_context, Arc::new(expr));
633629
expr_fn.elem.expr = Arc::unwrap_or_clone(inferred_expr);
634630
}
635631

@@ -1562,10 +1558,6 @@ impl ExprFunction {
15621558
annotate_variable(target, r#type.clone(), body)
15631559
});
15641560
let res = Expr::Lambda(*arity, Arc::new(new_body), meta.clone());
1565-
eprintln!(
1566-
"ASSIGN_PARAM_TYPES_TO_BODY_VARIABLES input:\n{:?}\nresult:\n{:?}",
1567-
self.expr, res
1568-
);
15691561
res
15701562
}
15711563
// TODO: Handle other cases - traverse the tree.
@@ -2141,6 +2133,13 @@ pub fn initial_context(ir: &IntermediateRepr) -> HashMap<Name, Expr<ExprMetadata
21412133
ctx
21422134
}
21432135

2136+
pub fn initial_typing_context(ir: &IntermediateRepr) -> HashMap<Name, FieldType> {
2137+
let ctx = initial_context(ir);
2138+
ctx.into_iter()
2139+
.filter_map(|(name, expr)| expr.meta().1.as_ref().map(|t| (name, t.clone())))
2140+
.collect()
2141+
}
2142+
21442143
#[cfg(test)]
21452144
mod tests {
21462145
use super::*;

‎engine/baml-lib/baml-core/src/validate/validation_pipeline/validations/expr_typecheck.rs

+69-87
Original file line numberDiff line numberDiff line change
@@ -157,95 +157,18 @@ pub fn typecheck_in_context(
157157
}
158158
}
159159
x => {
160-
eprintln!("TYPECHECKING APP: UNEXPECTED ARGS: {:?} {:?}", f, x);
160+
eprintln!(
161+
"TYPECHECKING APP: UNEXPECTED ARGS: ({}: {:?} ) {:?}",
162+
f.dump_str(),
163+
f.meta()
164+
.1
165+
.as_ref()
166+
.map_or("?".to_string(), |t| t.to_string()),
167+
x
168+
);
161169
}
162170
}
163171
Ok(())
164-
165-
// TODO: What was this? Bring it back?
166-
// match (f.as_ref(), xs.as_ref(), maybe_app_type) {
167-
// (
168-
// _, // Expr::Lambda(params, body, (lambda_span, _)),
169-
// Expr::ArgsTuple(args, (args_span, args_type)),
170-
// _,
171-
// ) => {
172-
// // First, check that the arguments are the right type
173-
// // for the lambda.
174-
// let maybe_lambda_type= &f.meta().1;
175-
// eprintln!("LAMBDA_TYPE: {:?}", maybe_lambda_type);
176-
// if let Some(lambda_type) = maybe_lambda_type {
177-
// eprintln!("checking lambda_type: {:?}", lambda_type);
178-
// match lambda_type {
179-
// ExprType::Arrow(arrow) => {
180-
// if let Some(app_type) = maybe_app_type {
181-
182-
// if !compatible_as_subtype(
183-
// ir,
184-
// &maybe_app_type,
185-
// &Some(arrow.body_type.clone()),
186-
// ) {
187-
// eprintln!(
188-
// "C Type mismatch in app: {} vs {}",
189-
// app_type.dump_str(),
190-
// arrow.body_type.dump_str()
191-
// );
192-
// diagnostics.push_error(DatamodelError::new_validation_error(
193-
// &format!(
194-
// "D Type mismatch in app: {} vs {}",
195-
// app_type.dump_str(),
196-
// arrow.body_type.dump_str()
197-
// ),
198-
// span.clone(),
199-
// ));
200-
// }
201-
// }
202-
// for (param_type, arg) in arrow.param_types.iter().zip(args.iter()) {
203-
// eprintln!("TYPECHECKING APP COMPARING PARAMTYPE: {:?} vs ARG: {:?}", param_type, arg);
204-
// if !compatible_as_subtype(
205-
// ir,
206-
// &arg.meta().1,
207-
// &Some(param_type.clone()),
208-
// ) {
209-
// eprintln!(
210-
// "E Type mismatch in app: {} vs {}",
211-
// arg.meta()
212-
// .1
213-
// .as_ref()
214-
// .map_or("?".to_string(), |t| t.dump_str()),
215-
// param_type.dump_str()
216-
// );
217-
// diagnostics.push_error(
218-
// DatamodelError::new_validation_error(
219-
// &format!(
220-
// "F Type mismatch in app: {} vs {}",
221-
// arg.meta()
222-
// .1
223-
// .as_ref()
224-
// .map_or("?".to_string(), |t| t.dump_str()),
225-
// param_type.dump_str()
226-
// ),
227-
// span.clone(),
228-
// ),
229-
// );
230-
// }
231-
// }
232-
// }
233-
// ExprType::Atom(_) => {
234-
// diagnostics.push_error(DatamodelError::new_validation_error(
235-
// "Expected a function type",
236-
// span.clone(),
237-
// ));
238-
// }
239-
// }
240-
// }
241-
242-
// typecheck_in_context(ir, diagnostics, &inner_context, body)?;
243-
244-
// Ok(())
245-
// }
246-
// _ => Ok(()),
247-
// }
248-
// Applications typecheck if the function arguments
249172
}
250173
Expr::Let(let_expr, _, _, _) => Ok(()),
251174
Expr::ArgsTuple(args, _) => Ok(()),
@@ -458,7 +381,66 @@ pub fn infer_types_in_context(
458381
(span.clone(), maybe_type.clone()),
459382
))
460383
}
461-
_ => expr.clone(),
384+
Expr::List(items, (span, maybe_type)) => {
385+
let new_items = items
386+
.iter()
387+
.map(|item| {
388+
Arc::unwrap_or_clone(infer_types_in_context(
389+
typing_context,
390+
Arc::new(item.clone()),
391+
))
392+
})
393+
.collect();
394+
Arc::new(Expr::List(new_items, (span.clone(), maybe_type.clone())))
395+
}
396+
Expr::Map(items, (span, maybe_type)) => {
397+
let new_items = items
398+
.iter()
399+
.map(|(key, value)| {
400+
(
401+
key.clone(),
402+
Arc::unwrap_or_clone(infer_types_in_context(
403+
typing_context,
404+
Arc::new(value.clone()),
405+
)),
406+
)
407+
})
408+
.collect();
409+
Arc::new(Expr::Map(new_items, (span.clone(), maybe_type.clone())))
410+
}
411+
Expr::ClassConstructor {
412+
name,
413+
fields,
414+
spread,
415+
meta,
416+
} => {
417+
let new_fields = fields
418+
.iter()
419+
.map(|(name, value)| {
420+
(
421+
name.clone(),
422+
Arc::unwrap_or_clone(infer_types_in_context(
423+
typing_context,
424+
Arc::new(value.clone()),
425+
)),
426+
)
427+
})
428+
.collect();
429+
let new_spread = spread.as_ref().map(|s| {
430+
Box::new(Arc::unwrap_or_clone(infer_types_in_context(
431+
typing_context,
432+
Arc::new(s.as_ref().clone()),
433+
)))
434+
});
435+
Arc::new(Expr::ClassConstructor {
436+
name: name.clone(),
437+
fields: new_fields,
438+
spread: new_spread,
439+
meta: meta.clone(),
440+
})
441+
}
442+
Expr::LLMFunction(llm_function, args, (span, maybe_type)) => expr.clone(),
443+
Expr::BoundVar(_, _) => expr.clone(),
462444
}
463445
}
464446

0 commit comments

Comments
 (0)