From efcb9e767a5313e2532023a28e37d97b67dde21d Mon Sep 17 00:00:00 2001 From: hellovai Date: Tue, 21 Oct 2025 23:31:57 -0700 Subject: [PATCH] Working block enter notifications --- engine/baml-compiler/src/codegen.rs | 31 + engine/baml-compiler/src/hir.rs | 6 + engine/baml-compiler/src/hir/dump.rs | 12 + engine/baml-compiler/src/hir/lowering.rs | 78 +- engine/baml-compiler/src/thir.rs | 23 + engine/baml-compiler/src/thir/interpret.rs | 1376 +++++++++-------- engine/baml-compiler/src/thir/typecheck.rs | 15 + engine/baml-compiler/src/watch.rs | 8 + engine/baml-compiler/src/watch/watch_event.rs | 6 +- engine/baml-runtime/src/async_vm_runtime.rs | 83 +- engine/baml-vm/src/bytecode.rs | 37 + engine/baml-vm/src/debug.rs | 4 + engine/baml-vm/src/vm.rs | 29 +- engine/baml-vm/tests/common.rs | 11 +- .../workflows/workflow_emit_simple.baml | 18 +- 15 files changed, 1023 insertions(+), 714 deletions(-) diff --git a/engine/baml-compiler/src/codegen.rs b/engine/baml-compiler/src/codegen.rs index 87a40d6a6e..3c00945e4c 100644 --- a/engine/baml-compiler/src/codegen.rs +++ b/engine/baml-compiler/src/codegen.rs @@ -551,6 +551,14 @@ impl<'g> HirCompiler<'g> { /// A statement is anything that does not produce a value by itself. fn compile_statement(&mut self, statement: &thir::Statement<(Span, Option)>) { match statement { + thir::Statement::AnnotatedStatement { headers, statement } => { + for header in headers { + self.emit_annotated_block(header); + } + if let Some(statement) = statement { + self.compile_statement(statement); + } + } thir::Statement::Let { name, value, .. } => { self.compile_expression(value); self.track_local(name); @@ -1513,6 +1521,29 @@ impl<'g> HirCompiler<'g> { self.emit(Instruction::LoadConst(const_index)); } + fn emit_annotated_block(&mut self, v: &str) { + self.emit_string_literal(v); + let mut function_name: [u8; 1024] = [0; 1024]; + function_name.copy_from_slice(v.as_bytes()); + // null terminate the vec in case its too long + function_name[std::cmp::min(v.len(), 1023)] = 0; + + let mut block_name: [u8; 1024] = [0; 1024]; + block_name.copy_from_slice(v.as_bytes()); + // null terminate the vec in case its too long + block_name[std::cmp::min(v.len(), 1023)] = 0; + + self.emit(Instruction::NotifyBlock( + baml_vm::bytecode::BlockNotification { + function_name, + block_name, + level: 1, + block_type: baml_vm::bytecode::BlockNotificationType::Statement, + is_enter: true, + }, + )); + } + /// Emits a single instruction and returns the index of the instruction. /// /// The return value is useful when we want to modify an instruction that diff --git a/engine/baml-compiler/src/hir.rs b/engine/baml-compiler/src/hir.rs index 683c7e23a5..3e6ee36648 100644 --- a/engine/baml-compiler/src/hir.rs +++ b/engine/baml-compiler/src/hir.rs @@ -201,6 +201,12 @@ pub enum Statement { variable: String, span: Span, }, + + /// Annotations that apply to the statement. + AnnotatedStatement { + headers: Vec, + statement: Option>, + }, } #[derive(Clone, Debug)] diff --git a/engine/baml-compiler/src/hir/dump.rs b/engine/baml-compiler/src/hir/dump.rs index ad5881ffbb..a88fe11587 100644 --- a/engine/baml-compiler/src/hir/dump.rs +++ b/engine/baml-compiler/src/hir/dump.rs @@ -132,6 +132,18 @@ impl TypeDocumentRender for TypeIR { impl Statement { pub fn to_doc(&self) -> RcDoc<'static, ()> { match self { + Statement::AnnotatedStatement { headers, statement } => { + let mut doc = RcDoc::text("//#") + .append(RcDoc::intersperse( + headers.iter().map(|h| RcDoc::text(h.clone())), + RcDoc::text(" "), + )) + .append(RcDoc::text("#//")); + if let Some(statement) = statement { + doc = doc.append(statement.to_doc().nest(2)); + } + doc + } Statement::Let { name, value, diff --git a/engine/baml-compiler/src/hir/lowering.rs b/engine/baml-compiler/src/hir/lowering.rs index cc4eef5578..91b25ae52b 100644 --- a/engine/baml-compiler/src/hir/lowering.rs +++ b/engine/baml-compiler/src/hir/lowering.rs @@ -367,19 +367,35 @@ impl Block { } // Second pass: lower statements, applying watch options to watch specs - let statements: Vec = block + let mut statements: Vec = block .stmts .iter() .map(|stmt| lower_stmt_with_options(stmt, &watch_options_map)) .collect(); + let trailing_expr = block + .expr + .as_deref() + .map(Expression::from_ast) + .map(Box::new); + + if !block.expr_headers.is_empty() { + println!( + "Annotated!: {}", + trailing_expr + .as_ref() + .map(|f| f.to_doc().pretty(80).to_string()) + .unwrap_or_else(|| "<..>".to_string()) + ); + statements.push(Statement::AnnotatedStatement { + headers: block.expr_headers.iter().map(|h| h.title.clone()).collect(), + statement: None, + }); + } + Block { statements, - trailing_expr: block - .expr - .as_deref() - .map(Expression::from_ast) - .map(Box::new), + trailing_expr, } } } @@ -388,6 +404,24 @@ fn lower_stmt(stmt: &ast::Stmt) -> Statement { lower_stmt_with_options(stmt, &HashMap::new()) } +fn maybe_annotated_statement( + stmt: Statement, + annotated_comments: &Vec>, +) -> Statement { + if annotated_comments.is_empty() { + stmt + } else { + println!("Annotated!: {}", stmt.to_doc().pretty(80)); + Statement::AnnotatedStatement { + headers: annotated_comments + .iter() + .map(|a| a.title.to_string()) + .collect(), + statement: Some(Box::new(stmt)), + } + } +} + fn lower_stmt_with_options( stmt: &ast::Stmt, watch_options: &HashMap, Option)>, @@ -489,7 +523,7 @@ fn lower_stmt_with_options( annotation, expr, span, - annotations: _, + annotations: annotated_comments, is_watched, }) => { let lifted_expr = Expression::from_ast(expr); @@ -504,7 +538,7 @@ fn lower_stmt_with_options( None }; - if *is_mutable { + let statement = if *is_mutable { Statement::DeclareAndAssign { name: identifier.to_string(), value: lifted_expr, @@ -520,7 +554,9 @@ fn lower_stmt_with_options( watch: watch_spec, span: span.clone(), } - } + }; + + maybe_annotated_statement(statement, annotated_comments) } ast::Stmt::ForLoop(ast::ForLoopStmt { identifier, @@ -528,23 +564,32 @@ fn lower_stmt_with_options( body, span, has_let: _, - annotations: _, + annotations: annotated_comments, }) => { // Lower for loop to HIR let lifted_iterator = Expression::from_ast(iterator); // Add the for loop statement - Statement::ForLoop { + let statement = Statement::ForLoop { identifier: identifier.name().to_string(), iterator: Box::new(lifted_iterator), block: Block::from_expr_block(body), span: span.clone(), - } + }; + + maybe_annotated_statement(statement, &annotated_comments) + } + ast::Stmt::Expression(ast::ExprStmt { + expr, + span, + annotations: annotated_comments, + }) => { + let statement = Statement::Expression { + expr: Expression::from_ast(expr), + span: span.clone(), + }; + maybe_annotated_statement(statement, &annotated_comments) } - ast::Stmt::Expression(expr) => Statement::Expression { - expr: Expression::from_ast(&expr.expr), - span: expr.span.clone(), - }, ast::Stmt::Semicolon(expr) => Statement::Semicolon { expr: Expression::from_ast(expr), span: expr.span().clone(), @@ -653,7 +698,6 @@ impl Expression { type_args, args, span, - .. }) => { // Note: AST function calls are always just names next to argument lists. // Later, we will be able to call any expression that is a function. diff --git a/engine/baml-compiler/src/thir.rs b/engine/baml-compiler/src/thir.rs index 926a67a9ee..a36dc6c4fa 100644 --- a/engine/baml-compiler/src/thir.rs +++ b/engine/baml-compiler/src/thir.rs @@ -720,6 +720,12 @@ pub enum Statement { variable: String, span: Span, }, + + /// Annotations that apply to the statement. + AnnotatedStatement { + headers: Vec, + statement: Option>>, + }, } impl Statement { @@ -728,6 +734,19 @@ impl Statement { T: std::fmt::Debug, { match self { + Statement::AnnotatedStatement { headers, statement } => { + let headers_str = + headers + .iter() + .map(|h| format!("//# {h}")) + .chain(std::iter::once( + statement + .as_ref() + .map(|s| s.dump_str()) + .unwrap_or_else(String::new), + )); + join(headers_str, "\n") + } Statement::Let { name, value, @@ -838,6 +857,10 @@ impl Statement { T: Clone, { match self { + Statement::AnnotatedStatement { statement, .. } => statement + .as_ref() + .map(|s| s.variables()) + .unwrap_or_else(HashSet::new), Statement::Declare { .. } | Statement::Break(_) | Statement::Continue(_) => { HashSet::new() } diff --git a/engine/baml-compiler/src/thir/interpret.rs b/engine/baml-compiler/src/thir/interpret.rs index 86e9b212f8..fd51c70961 100644 --- a/engine/baml-compiler/src/thir/interpret.rs +++ b/engine/baml-compiler/src/thir/interpret.rs @@ -513,62 +513,44 @@ where block.statements.len() }; - for stmt in block.statements.iter().take(statements_to_execute) { - match stmt { - Statement::Let { - name, value, watch, .. - } => { - match evaluate_expr( - value, - scopes, - thir, - run_llm_function, - watch_handler, - function_name, - ) - .await? - { - EvalValue::Value(v) => { - declare(scopes, name, v); - // Register watch tracking if @watch is present - if let Some(watch_spec) = watch { - if let Some(var_ref) = lookup_variable(scopes, name) { - register_watch_variable( - scopes, - name, - var_ref, - watch_spec.clone(), - ); - } - } - } - EvalValue::Reference(cell) => { - declare_with_cell(scopes, name, cell.clone()); - // Register watch tracking if @watch is present - if let Some(emit_spec) = watch { - register_watch_variable(scopes, name, cell, emit_spec.clone()); - } - } - EvalValue::Function(_, _, _) => { - bail!("cannot assign function to variable `{}`", name); + fn handle_statement<'a, F, Fut, E>( + stmt: &'a Statement, + scopes: &'a mut Vec, + thir: &'a THir, + run_llm_function: &'a mut F, + watch_handler: &'a mut E, + function_name: &'a str, + ) -> BoxFuture<'a, Result>> + where + F: LlmHandler, + Fut: LlmFuture, + E: FnMut(crate::watch::WatchNotification) + Send, + { + Box::pin(async move { + match stmt { + Statement::AnnotatedStatement { headers, statement } => { + headers.iter().for_each(|header| { + watch_handler(WatchNotification::new_block( + header.clone(), + "FunctionNameNotPlumbed".to_string(), + )); + }); + if let Some(statement) = statement { + return handle_statement( + statement, + scopes, + thir, + run_llm_function, + watch_handler, + function_name, + ) + .await; } } - // Check for changes in all watch variables after the let statement - check_watch_changes( - scopes, - watch_handler, - function_name, - thir, - run_llm_function, - ) - .await; - } - Statement::Declare { name, span } => { - declare(scopes, name, BamlValueWithMeta::Null((span.clone(), None))); - } - Statement::Assign { left, value } => { - let assigned_value = expect_value( - evaluate_expr( + Statement::Let { + name, value, watch, .. + } => { + match evaluate_expr( value, scopes, thir, @@ -576,191 +558,192 @@ where watch_handler, function_name, ) - .await?, - )?; - assign_to_expr( - left, - assigned_value, - scopes, - thir, - run_llm_function, - watch_handler, - function_name, - ) - .await?; - // Check for changes in watch variables after assignment - check_watch_changes( - scopes, - watch_handler, - function_name, - thir, - run_llm_function, - ) - .await; - } - Statement::DeclareAndAssign { - name, value, watch, .. - } => { - // Create watch context if @watch is present - let watch_ctx = watch.as_ref().map(|_| { - use std::time::{SystemTime, UNIX_EPOCH}; - let timestamp = SystemTime::now() - .duration_since(UNIX_EPOCH) - .unwrap() - .as_millis(); - WatchStreamContext { - variable_name: name.clone(), - stream_id: format!("{function_name}_{name}_{timestamp}"), - } - }); - - match evaluate_expr_with_context( - value, - scopes, - thir, - run_llm_function, - watch_handler, - function_name, - watch_ctx.as_ref(), - ) - .await? - { - EvalValue::Value(v) => { - declare(scopes, name, v); - // Register watch tracking if @watch is present - if let Some(watch_spec) = watch { - if let Some(var_ref) = lookup_variable(scopes, name) { - register_watch_variable( - scopes, - name, - var_ref, - watch_spec.clone(), - ); + .await? + { + EvalValue::Value(v) => { + declare(scopes, name, v); + // Register watch tracking if @watch is present + if let Some(watch_spec) = watch { + if let Some(var_ref) = lookup_variable(scopes, name) { + register_watch_variable( + scopes, + name, + var_ref, + watch_spec.clone(), + ); + } } } - } - EvalValue::Reference(cell) => { - declare_with_cell(scopes, name, cell.clone()); - // Register watch tracking if @watch is present - if let Some(emit_spec) = watch { - register_watch_variable(scopes, name, cell, emit_spec.clone()); + EvalValue::Reference(cell) => { + declare_with_cell(scopes, name, cell.clone()); + // Register watch tracking if @watch is present + if let Some(emit_spec) = watch { + register_watch_variable(scopes, name, cell, emit_spec.clone()); + } + } + EvalValue::Function(_, _, _) => { + bail!("cannot assign function to variable `{}`", name); } } - EvalValue::Function(_, _, _) => { - bail!("cannot assign function to variable `{}`", name); - } + // Check for changes in all watch variables after the let statement + check_watch_changes( + scopes, + watch_handler, + function_name, + thir, + run_llm_function, + ) + .await; } - // Check for changes in all watch variables after the declare and assign - check_watch_changes( - scopes, - watch_handler, - function_name, - thir, - run_llm_function, - ) - .await; - } - Statement::Return { expr, .. } => { - let v = expect_value( - evaluate_expr( - expr, + Statement::Declare { name, span } => { + declare(scopes, name, BamlValueWithMeta::Null((span.clone(), None))); + } + Statement::Assign { left, value } => { + let assigned_value = expect_value( + evaluate_expr( + value, + scopes, + thir, + run_llm_function, + watch_handler, + function_name, + ) + .await?, + )?; + assign_to_expr( + left, + assigned_value, scopes, thir, run_llm_function, watch_handler, function_name, ) - .await?, - )?; - scopes.pop(); - return Ok(ControlFlow::Return(v)); - } - Statement::Expression { expr, .. } => { - // For expression statements, we still need to evaluate them for side effects - // (and the last one might be the implicit return value) - let _ = evaluate_expr( - expr, - scopes, - thir, - run_llm_function, - watch_handler, - function_name, - ) - .await?; - } - Statement::Break(_) => { - scopes.pop(); - return Ok(ControlFlow::Break); - } - Statement::Continue(_) => { - scopes.pop(); - return Ok(ControlFlow::Continue); - } - Statement::While { - condition, block, .. - } => loop { - let cond_val = expect_value( - evaluate_expr( - condition, + .await?; + // Check for changes in watch variables after assignment + check_watch_changes( scopes, - thir, - run_llm_function, watch_handler, function_name, + thir, + run_llm_function, ) - .await?, - )?; - match cond_val { - BamlValueWithMeta::Bool(true, _) => match evaluate_block_with_control_flow( - block, + .await; + } + Statement::DeclareAndAssign { + name, value, watch, .. + } => { + // Create watch context if @watch is present + let watch_ctx = watch.as_ref().map(|_| { + use std::time::{SystemTime, UNIX_EPOCH}; + let timestamp = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_millis(); + WatchStreamContext { + variable_name: name.clone(), + stream_id: format!("{function_name}_{name}_{timestamp}"), + } + }); + + match evaluate_expr_with_context( + value, scopes, thir, run_llm_function, watch_handler, function_name, + watch_ctx.as_ref(), ) .await? { - ControlFlow::Break => break, - ControlFlow::Continue => continue, - ControlFlow::Normal(_) => {} - ControlFlow::Return(val) => { - scopes.pop(); - return Ok(ControlFlow::Return(val)); + EvalValue::Value(v) => { + declare(scopes, name, v); + // Register watch tracking if @watch is present + if let Some(watch_spec) = watch { + if let Some(var_ref) = lookup_variable(scopes, name) { + register_watch_variable( + scopes, + name, + var_ref, + watch_spec.clone(), + ); + } + } + } + EvalValue::Reference(cell) => { + declare_with_cell(scopes, name, cell.clone()); + // Register watch tracking if @watch is present + if let Some(emit_spec) = watch { + register_watch_variable(scopes, name, cell, emit_spec.clone()); + } + } + EvalValue::Function(_, _, _) => { + bail!("cannot assign function to variable `{}`", name); } - }, - BamlValueWithMeta::Bool(false, _) => break, - _ => bail!("while condition must be boolean"), + } + // Check for changes in all watch variables after the declare and assign + check_watch_changes( + scopes, + watch_handler, + function_name, + thir, + run_llm_function, + ) + .await; } - }, - Statement::ForLoop { - identifier, - iterator, - block, - .. - } => { - let iterable_val = expect_value( - evaluate_expr( - iterator, + Statement::Return { expr, .. } => { + let v = expect_value( + evaluate_expr( + expr, + scopes, + thir, + run_llm_function, + watch_handler, + function_name, + ) + .await?, + )?; + scopes.pop(); + return Ok(Some(ControlFlow::Return(v))); + } + Statement::Expression { expr, .. } => { + // For expression statements, we still need to evaluate them for side effects + // (and the last one might be the implicit return value) + let _ = evaluate_expr( + expr, scopes, thir, run_llm_function, watch_handler, function_name, ) - .await?, - )?; - match iterable_val { - BamlValueWithMeta::List(items, _) => { - for item_val in items.iter() { - // Create new scope for loop iteration - scopes.push(Scope { - variables: BamlMap::new(), - watch_variables: Vec::new(), - is_filter_context: false, - }); - declare(scopes, identifier, item_val.clone()); - + .await?; + } + Statement::Break(_) => { + scopes.pop(); + return Ok(Some(ControlFlow::Break)); + } + Statement::Continue(_) => { + scopes.pop(); + return Ok(Some(ControlFlow::Continue)); + } + Statement::While { + condition, block, .. + } => loop { + let cond_val = expect_value( + evaluate_expr( + condition, + scopes, + thir, + run_llm_function, + watch_handler, + function_name, + ) + .await?, + )?; + match cond_val { + BamlValueWithMeta::Bool(true, _) => { match evaluate_block_with_control_flow( block, scopes, @@ -771,516 +754,609 @@ where ) .await? { - ControlFlow::Break => { - scopes.pop(); - break; - } - ControlFlow::Continue => { - scopes.pop(); - continue; - } - ControlFlow::Normal(_) => { - scopes.pop(); - } + ControlFlow::Break => break, + ControlFlow::Continue => continue, + ControlFlow::Normal(_) => {} ControlFlow::Return(val) => { scopes.pop(); - scopes.pop(); - return Ok(ControlFlow::Return(val)); + return Ok(Some(ControlFlow::Return(val))); + } + } + } + BamlValueWithMeta::Bool(false, _) => break, + _ => bail!("while condition must be boolean"), + } + }, + Statement::ForLoop { + identifier, + iterator, + block, + .. + } => { + let iterable_val = expect_value( + evaluate_expr( + iterator, + scopes, + thir, + run_llm_function, + watch_handler, + function_name, + ) + .await?, + )?; + match iterable_val { + BamlValueWithMeta::List(items, _) => { + for item_val in items.iter() { + // Create new scope for loop iteration + scopes.push(Scope { + variables: BamlMap::new(), + watch_variables: Vec::new(), + is_filter_context: false, + }); + declare(scopes, identifier, item_val.clone()); + + match evaluate_block_with_control_flow( + block, + scopes, + thir, + run_llm_function, + watch_handler, + function_name, + ) + .await? + { + ControlFlow::Break => { + scopes.pop(); + return Ok(Some(ControlFlow::Break)); + } + ControlFlow::Continue => { + scopes.pop(); + continue; + } + ControlFlow::Normal(_) => { + scopes.pop(); + } + ControlFlow::Return(val) => { + scopes.pop(); + scopes.pop(); + return Ok(Some(ControlFlow::Return(val))); + } } } } + _ => bail!("for loop requires iterable (list)"), } - _ => bail!("for loop requires iterable (list)"), } - } - Statement::AssignOp { - left, - value, - assign_op, - .. - } => { - use crate::hir::AssignOp; - - let current_val = expect_value( - evaluate_expr( - left, - scopes, - thir, - run_llm_function, - watch_handler, - function_name, - ) - .await?, - )?; + Statement::AssignOp { + left, + value, + assign_op, + .. + } => { + use crate::hir::AssignOp; - // Evaluate the right-hand side expression - let rhs_val = expect_value( - evaluate_expr( - value, - scopes, - thir, - run_llm_function, - watch_handler, - function_name, - ) - .await?, - )?; + let current_val = expect_value( + evaluate_expr( + left, + scopes, + thir, + run_llm_function, + watch_handler, + function_name, + ) + .await?, + )?; - // Perform the compound assignment operation - let result_val = match assign_op { - AssignOp::AddAssign => match (current_val.clone(), rhs_val.clone()) { - (BamlValueWithMeta::Int(a, meta), BamlValueWithMeta::Int(b, _)) => { - BamlValueWithMeta::Int(a + b, meta) - } - (BamlValueWithMeta::Float(a, meta), BamlValueWithMeta::Float(b, _)) => { - BamlValueWithMeta::Float(a + b, meta) - } - (BamlValueWithMeta::Int(a, meta), BamlValueWithMeta::Float(b, _)) => { - BamlValueWithMeta::Float(a as f64 + b, meta) - } - (BamlValueWithMeta::Float(a, meta), BamlValueWithMeta::Int(b, _)) => { - BamlValueWithMeta::Float(a + (b as f64), meta) - } - ( - BamlValueWithMeta::String(a, meta), - BamlValueWithMeta::String(b, _), - ) => BamlValueWithMeta::String(format!("{a}{b}"), meta), - _ => bail!("unsupported types for += operator"), - }, - AssignOp::SubAssign => match (current_val.clone(), rhs_val.clone()) { - (BamlValueWithMeta::Int(a, meta), BamlValueWithMeta::Int(b, _)) => { - BamlValueWithMeta::Int(a - b, meta) - } - (BamlValueWithMeta::Float(a, meta), BamlValueWithMeta::Float(b, _)) => { - BamlValueWithMeta::Float(a - b, meta) - } - (BamlValueWithMeta::Int(a, meta), BamlValueWithMeta::Float(b, _)) => { - BamlValueWithMeta::Float((a as f64) - b, meta) - } - (BamlValueWithMeta::Float(a, meta), BamlValueWithMeta::Int(b, _)) => { - BamlValueWithMeta::Float(a - (b as f64), meta) - } - _ => bail!("unsupported types for -= operator"), - }, - AssignOp::MulAssign => match (current_val.clone(), rhs_val.clone()) { - (BamlValueWithMeta::Int(a, meta), BamlValueWithMeta::Int(b, _)) => { - BamlValueWithMeta::Int(a * b, meta) - } - (BamlValueWithMeta::Float(a, meta), BamlValueWithMeta::Float(b, _)) => { - BamlValueWithMeta::Float(a * b, meta) - } - (BamlValueWithMeta::Int(a, meta), BamlValueWithMeta::Float(b, _)) => { - BamlValueWithMeta::Float((a as f64) * b, meta) - } - (BamlValueWithMeta::Float(a, meta), BamlValueWithMeta::Int(b, _)) => { - BamlValueWithMeta::Float(a * (b as f64), meta) - } - _ => bail!("unsupported types for *= operator"), - }, - AssignOp::DivAssign => match (current_val.clone(), rhs_val.clone()) { - (BamlValueWithMeta::Int(a, meta), BamlValueWithMeta::Int(b, _)) => { - if b == 0 { - bail!("division by zero in /= operator"); + // Evaluate the right-hand side expression + let rhs_val = expect_value( + evaluate_expr( + value, + scopes, + thir, + run_llm_function, + watch_handler, + function_name, + ) + .await?, + )?; + + // Perform the compound assignment operation + let result_val = match assign_op { + AssignOp::AddAssign => match (current_val.clone(), rhs_val.clone()) { + (BamlValueWithMeta::Int(a, meta), BamlValueWithMeta::Int(b, _)) => { + BamlValueWithMeta::Int(a + b, meta) } - BamlValueWithMeta::Float((a as f64) / (b as f64), meta) - } - (BamlValueWithMeta::Float(a, meta), BamlValueWithMeta::Float(b, _)) => { - if b == 0.0 { - bail!("division by zero in /= operator"); + ( + BamlValueWithMeta::Float(a, meta), + BamlValueWithMeta::Float(b, _), + ) => BamlValueWithMeta::Float(a + b, meta), + ( + BamlValueWithMeta::Int(a, meta), + BamlValueWithMeta::Float(b, _), + ) => BamlValueWithMeta::Float(a as f64 + b, meta), + ( + BamlValueWithMeta::Float(a, meta), + BamlValueWithMeta::Int(b, _), + ) => BamlValueWithMeta::Float(a + (b as f64), meta), + ( + BamlValueWithMeta::String(a, meta), + BamlValueWithMeta::String(b, _), + ) => BamlValueWithMeta::String(format!("{a}{b}"), meta), + _ => bail!("unsupported types for += operator"), + }, + AssignOp::SubAssign => match (current_val.clone(), rhs_val.clone()) { + (BamlValueWithMeta::Int(a, meta), BamlValueWithMeta::Int(b, _)) => { + BamlValueWithMeta::Int(a - b, meta) } - BamlValueWithMeta::Float(a / b, meta) - } - (BamlValueWithMeta::Int(a, meta), BamlValueWithMeta::Float(b, _)) => { - if b == 0.0 { - bail!("division by zero in /= operator"); + ( + BamlValueWithMeta::Float(a, meta), + BamlValueWithMeta::Float(b, _), + ) => BamlValueWithMeta::Float(a - b, meta), + ( + BamlValueWithMeta::Int(a, meta), + BamlValueWithMeta::Float(b, _), + ) => BamlValueWithMeta::Float((a as f64) - b, meta), + ( + BamlValueWithMeta::Float(a, meta), + BamlValueWithMeta::Int(b, _), + ) => BamlValueWithMeta::Float(a - (b as f64), meta), + _ => bail!("unsupported types for -= operator"), + }, + AssignOp::MulAssign => match (current_val.clone(), rhs_val.clone()) { + (BamlValueWithMeta::Int(a, meta), BamlValueWithMeta::Int(b, _)) => { + BamlValueWithMeta::Int(a * b, meta) } - BamlValueWithMeta::Float((a as f64) / b, meta) - } - (BamlValueWithMeta::Float(a, meta), BamlValueWithMeta::Int(b, _)) => { - if b == 0 { - bail!("division by zero in /= operator"); + ( + BamlValueWithMeta::Float(a, meta), + BamlValueWithMeta::Float(b, _), + ) => BamlValueWithMeta::Float(a * b, meta), + ( + BamlValueWithMeta::Int(a, meta), + BamlValueWithMeta::Float(b, _), + ) => BamlValueWithMeta::Float((a as f64) * b, meta), + ( + BamlValueWithMeta::Float(a, meta), + BamlValueWithMeta::Int(b, _), + ) => BamlValueWithMeta::Float(a * (b as f64), meta), + _ => bail!("unsupported types for *= operator"), + }, + AssignOp::DivAssign => match (current_val.clone(), rhs_val.clone()) { + (BamlValueWithMeta::Int(a, meta), BamlValueWithMeta::Int(b, _)) => { + if b == 0 { + bail!("division by zero in /= operator"); + } + BamlValueWithMeta::Float((a as f64) / (b as f64), meta) } - BamlValueWithMeta::Float(a / (b as f64), meta) - } - _ => bail!("unsupported types for /= operator"), - }, - AssignOp::ModAssign => match (current_val.clone(), rhs_val.clone()) { - (BamlValueWithMeta::Int(a, meta), BamlValueWithMeta::Int(b, _)) => { - if b == 0 { - bail!("modulo by zero in %= operator"); + ( + BamlValueWithMeta::Float(a, meta), + BamlValueWithMeta::Float(b, _), + ) => { + if b == 0.0 { + bail!("division by zero in /= operator"); + } + BamlValueWithMeta::Float(a / b, meta) } - BamlValueWithMeta::Int(a % b, meta) - } - _ => bail!("unsupported types for %= operator"), - }, - AssignOp::BitXorAssign => match (current_val.clone(), rhs_val.clone()) { - (BamlValueWithMeta::Int(a, meta), BamlValueWithMeta::Int(b, _)) => { - BamlValueWithMeta::Int(a ^ b, meta) - } - _ => bail!("bitwise ^= requires integer operands"), - }, - AssignOp::BitAndAssign => match (current_val.clone(), rhs_val.clone()) { - (BamlValueWithMeta::Int(a, meta), BamlValueWithMeta::Int(b, _)) => { - BamlValueWithMeta::Int(a & b, meta) - } - _ => bail!("bitwise &= requires integer operands"), - }, - AssignOp::BitOrAssign => match (current_val.clone(), rhs_val.clone()) { - (BamlValueWithMeta::Int(a, meta), BamlValueWithMeta::Int(b, _)) => { - BamlValueWithMeta::Int(a | b, meta) - } - _ => bail!("bitwise |= requires integer operands"), - }, - AssignOp::ShlAssign => match (current_val.clone(), rhs_val.clone()) { - (BamlValueWithMeta::Int(a, meta), BamlValueWithMeta::Int(b, _)) => { - if b < 0 { - bail!("negative shift amount in <<= operator"); + ( + BamlValueWithMeta::Int(a, meta), + BamlValueWithMeta::Float(b, _), + ) => { + if b == 0.0 { + bail!("division by zero in /= operator"); + } + BamlValueWithMeta::Float((a as f64) / b, meta) } - BamlValueWithMeta::Int(a << b, meta) - } - _ => bail!("shift <<= requires integer operands"), - }, - AssignOp::ShrAssign => match (current_val.clone(), rhs_val.clone()) { - (BamlValueWithMeta::Int(a, meta), BamlValueWithMeta::Int(b, _)) => { - if b < 0 { - bail!("negative shift amount in >>= operator"); + ( + BamlValueWithMeta::Float(a, meta), + BamlValueWithMeta::Int(b, _), + ) => { + if b == 0 { + bail!("division by zero in /= operator"); + } + BamlValueWithMeta::Float(a / (b as f64), meta) + } + _ => bail!("unsupported types for /= operator"), + }, + AssignOp::ModAssign => match (current_val.clone(), rhs_val.clone()) { + (BamlValueWithMeta::Int(a, meta), BamlValueWithMeta::Int(b, _)) => { + if b == 0 { + bail!("modulo by zero in %= operator"); + } + BamlValueWithMeta::Int(a % b, meta) + } + _ => bail!("unsupported types for %= operator"), + }, + AssignOp::BitXorAssign => { + match (current_val.clone(), rhs_val.clone()) { + ( + BamlValueWithMeta::Int(a, meta), + BamlValueWithMeta::Int(b, _), + ) => BamlValueWithMeta::Int(a ^ b, meta), + _ => bail!("bitwise ^= requires integer operands"), } - BamlValueWithMeta::Int(a >> b, meta) } - _ => bail!("shift >>= requires integer operands"), - }, - }; - - // Assign the result back to the target expression - assign_to_expr( - left, - result_val, - scopes, - thir, - run_llm_function, - watch_handler, - function_name, - ) - .await?; - // Check for changes in watch variables after compound assignment - check_watch_changes( - scopes, - watch_handler, - function_name, - thir, - run_llm_function, - ) - .await; - } - Statement::SemicolonExpression { expr, .. } => { - let _ = evaluate_expr( - expr, - scopes, - thir, - run_llm_function, - watch_handler, - function_name, - ) - .await?; - } - Statement::CForLoop { - condition, - after, - block, - } => { - loop { - // Check condition (if present) - if let Some(cond_expr) = condition { - let cond_val = expect_value( - evaluate_expr( - cond_expr, - scopes, - thir, - run_llm_function, - watch_handler, - function_name, - ) - .await?, - )?; - match cond_val { - BamlValueWithMeta::Bool(false, _) => break, - BamlValueWithMeta::Bool(true, _) => {} - _ => bail!("C-style for loop condition must be boolean"), + AssignOp::BitAndAssign => { + match (current_val.clone(), rhs_val.clone()) { + ( + BamlValueWithMeta::Int(a, meta), + BamlValueWithMeta::Int(b, _), + ) => BamlValueWithMeta::Int(a & b, meta), + _ => bail!("bitwise &= requires integer operands"), + } } - } + AssignOp::BitOrAssign => match (current_val.clone(), rhs_val.clone()) { + (BamlValueWithMeta::Int(a, meta), BamlValueWithMeta::Int(b, _)) => { + BamlValueWithMeta::Int(a | b, meta) + } + _ => bail!("bitwise |= requires integer operands"), + }, + AssignOp::ShlAssign => match (current_val.clone(), rhs_val.clone()) { + (BamlValueWithMeta::Int(a, meta), BamlValueWithMeta::Int(b, _)) => { + if b < 0 { + bail!("negative shift amount in <<= operator"); + } + BamlValueWithMeta::Int(a << b, meta) + } + _ => bail!("shift <<= requires integer operands"), + }, + AssignOp::ShrAssign => match (current_val.clone(), rhs_val.clone()) { + (BamlValueWithMeta::Int(a, meta), BamlValueWithMeta::Int(b, _)) => { + if b < 0 { + bail!("negative shift amount in >>= operator"); + } + BamlValueWithMeta::Int(a >> b, meta) + } + _ => bail!("shift >>= requires integer operands"), + }, + }; - // Execute loop body - match evaluate_block_with_control_flow( - block, + // Assign the result back to the target expression + assign_to_expr( + left, + result_val, scopes, thir, run_llm_function, watch_handler, function_name, ) - .await? - { - ControlFlow::Break => break, - ControlFlow::Continue => { - // Execute after statement if present - if let Some(after_stmt) = after { - // Execute the after statement in the current scope context - match after_stmt.as_ref() { - Statement::AssignOp { - left, - value, - assign_op, - .. - } => { - use crate::hir::AssignOp; - - let current_val = expect_value( - evaluate_expr( + .await?; + // Check for changes in watch variables after compound assignment + check_watch_changes( + scopes, + watch_handler, + function_name, + thir, + run_llm_function, + ) + .await; + } + Statement::SemicolonExpression { expr, .. } => { + let _ = evaluate_expr( + expr, + scopes, + thir, + run_llm_function, + watch_handler, + function_name, + ) + .await?; + } + Statement::CForLoop { + condition, + after, + block, + } => { + loop { + // Check condition (if present) + if let Some(cond_expr) = condition { + let cond_val = expect_value( + evaluate_expr( + cond_expr, + scopes, + thir, + run_llm_function, + watch_handler, + function_name, + ) + .await?, + )?; + match cond_val { + BamlValueWithMeta::Bool(false, _) => break, + BamlValueWithMeta::Bool(true, _) => {} + _ => bail!("C-style for loop condition must be boolean"), + } + } + + // Execute loop body + match evaluate_block_with_control_flow( + block, + scopes, + thir, + run_llm_function, + watch_handler, + function_name, + ) + .await? + { + ControlFlow::Break => break, + ControlFlow::Continue => { + // Execute after statement if present + if let Some(after_stmt) = after { + // Execute the after statement in the current scope context + match after_stmt.as_ref() { + Statement::AssignOp { + left, + value, + assign_op, + .. + } => { + use crate::hir::AssignOp; + + let current_val = expect_value( + evaluate_expr( + left, + scopes, + thir, + run_llm_function, + watch_handler, + function_name, + ) + .await?, + )?; + let rhs_val = expect_value( + evaluate_expr( + value, + scopes, + thir, + run_llm_function, + watch_handler, + function_name, + ) + .await?, + )?; + + let result_val = match assign_op { + AssignOp::AddAssign => { + match (current_val.clone(), rhs_val.clone()) { + ( + BamlValueWithMeta::Int(a, meta), + BamlValueWithMeta::Int(b, _), + ) => BamlValueWithMeta::Int(a + b, meta), + _ => bail!( + "unsupported types for += in C-for after clause" + ), + } + } + _ => bail!( + "unsupported assign op in C-for after clause" + ), + }; + assign_to_expr( left, + result_val, scopes, thir, run_llm_function, watch_handler, function_name, ) - .await?, - )?; - let rhs_val = expect_value( - evaluate_expr( - value, - scopes, - thir, - run_llm_function, - watch_handler, - function_name, - ) - .await?, - )?; - - let result_val = match assign_op { - AssignOp::AddAssign => { - match (current_val.clone(), rhs_val.clone()) { - ( - BamlValueWithMeta::Int(a, meta), - BamlValueWithMeta::Int(b, _), - ) => BamlValueWithMeta::Int(a + b, meta), - _ => bail!( - "unsupported types for += in C-for after clause" - ), - } - } - _ => bail!( - "unsupported assign op in C-for after clause" - ), - }; - assign_to_expr( - left, - result_val, - scopes, - thir, - run_llm_function, - watch_handler, - function_name, - ) - .await?; - } - Statement::Assign { left, value } => { - let v = expect_value( - evaluate_expr( - value, + .await?; + } + Statement::Assign { left, value } => { + let v = expect_value( + evaluate_expr( + value, + scopes, + thir, + run_llm_function, + watch_handler, + function_name, + ) + .await?, + )?; + assign_to_expr( + left, + v, scopes, thir, run_llm_function, watch_handler, function_name, ) - .await?, - )?; - assign_to_expr( - left, - v, - scopes, - thir, - run_llm_function, - watch_handler, - function_name, - ) - .await?; + .await?; + } + _ => bail!( + "unsupported statement type in C-for after clause" + ), } - _ => bail!( - "unsupported statement type in C-for after clause" - ), } + continue; } - continue; - } - ControlFlow::Normal(_) => { - // Execute after statement if present - if let Some(after_stmt) = after { - // Execute the after statement in the current scope context - match after_stmt.as_ref() { - Statement::AssignOp { - left, - value, - assign_op, - .. - } => { - use crate::hir::AssignOp; - - let current_val = expect_value( - evaluate_expr( + ControlFlow::Normal(_) => { + // Execute after statement if present + if let Some(after_stmt) = after { + // Execute the after statement in the current scope context + match after_stmt.as_ref() { + Statement::AssignOp { + left, + value, + assign_op, + .. + } => { + use crate::hir::AssignOp; + + let current_val = expect_value( + evaluate_expr( + left, + scopes, + thir, + run_llm_function, + watch_handler, + function_name, + ) + .await?, + )?; + let rhs_val = expect_value( + evaluate_expr( + value, + scopes, + thir, + run_llm_function, + watch_handler, + function_name, + ) + .await?, + )?; + + let result_val = match assign_op { + AssignOp::AddAssign => { + match (current_val.clone(), rhs_val.clone()) { + ( + BamlValueWithMeta::Int(a, meta), + BamlValueWithMeta::Int(b, _), + ) => BamlValueWithMeta::Int(a + b, meta), + _ => bail!( + "unsupported types for += in C-for after clause" + ), + } + } + _ => bail!( + "unsupported assign op in C-for after clause" + ), + }; + assign_to_expr( left, + result_val, scopes, thir, run_llm_function, watch_handler, function_name, ) - .await?, - )?; - let rhs_val = expect_value( - evaluate_expr( - value, - scopes, - thir, - run_llm_function, - watch_handler, - function_name, - ) - .await?, - )?; - - let result_val = match assign_op { - AssignOp::AddAssign => { - match (current_val.clone(), rhs_val.clone()) { - ( - BamlValueWithMeta::Int(a, meta), - BamlValueWithMeta::Int(b, _), - ) => BamlValueWithMeta::Int(a + b, meta), - _ => bail!( - "unsupported types for += in C-for after clause" - ), - } - } - _ => bail!( - "unsupported assign op in C-for after clause" - ), - }; - assign_to_expr( - left, - result_val, - scopes, - thir, - run_llm_function, - watch_handler, - function_name, - ) - .await?; - } - Statement::Assign { left, value } => { - let v = expect_value( - evaluate_expr( - value, + .await?; + } + Statement::Assign { left, value } => { + let v = expect_value( + evaluate_expr( + value, + scopes, + thir, + run_llm_function, + watch_handler, + function_name, + ) + .await?, + )?; + assign_to_expr( + left, + v, scopes, thir, run_llm_function, watch_handler, function_name, ) - .await?, - )?; - assign_to_expr( - left, - v, - scopes, - thir, - run_llm_function, - watch_handler, - function_name, - ) - .await?; + .await?; + } + _ => bail!( + "unsupported statement type in C-for after clause" + ), } - _ => bail!( - "unsupported statement type in C-for after clause" - ), } } - } - ControlFlow::Return(val) => { - scopes.pop(); - return Ok(ControlFlow::Return(val)); + ControlFlow::Return(val) => { + scopes.pop(); + return Ok(Some(ControlFlow::Return(val))); + } } } } - } - Statement::Assert { condition, .. } => { - let cond_val = expect_value( - evaluate_expr( - condition, - scopes, - thir, - run_llm_function, - watch_handler, - function_name, - ) - .await?, - )?; - match cond_val { - BamlValueWithMeta::Bool(true, _) => {} - BamlValueWithMeta::Bool(false, _) => bail!("assertion failed"), - _ => bail!("assert condition must be boolean"), + Statement::Assert { condition, .. } => { + let cond_val = expect_value( + evaluate_expr( + condition, + scopes, + thir, + run_llm_function, + watch_handler, + function_name, + ) + .await?, + )?; + match cond_val { + BamlValueWithMeta::Bool(true, _) => {} + BamlValueWithMeta::Bool(false, _) => bail!("assertion failed"), + _ => bail!("assert condition must be boolean"), + } } - } - Statement::WatchOptions { - variable, - channel, - when, - span, - } => { - // Find and update the watch variable for this variable - // We need to find the watch variable by checking which one references the same value - for scope in scopes.iter_mut().rev() { - if let Some(var_ref) = scope.variables.get(variable) { - // Find the watch variable that references this variable - if let Some(watch_var) = scope - .watch_variables - .iter_mut() - .find(|wv| Arc::ptr_eq(&wv.value_ref, var_ref)) - { - // Update the channel name if provided - if let Some(new_channel) = channel { - watch_var.spec.name = new_channel.clone(); - } + Statement::WatchOptions { + variable, + channel, + when, + span, + } => { + // Find and update the watch variable for this variable + // We need to find the watch variable by checking which one references the same value + for scope in scopes.iter_mut().rev() { + if let Some(var_ref) = scope.variables.get(variable) { + // Find the watch variable that references this variable + if let Some(watch_var) = scope + .watch_variables + .iter_mut() + .find(|wv| Arc::ptr_eq(&wv.value_ref, var_ref)) + { + // Update the channel name if provided + if let Some(new_channel) = channel { + watch_var.spec.name = new_channel.clone(); + } - // Update the when condition if provided - if let Some(when_str) = when { - watch_var.spec.when = match when_str.as_str() { - "manual" => crate::watch::WatchWhen::Manual, - "true" => crate::watch::WatchWhen::True, - _ => crate::watch::WatchWhen::FunctionName( - internal_baml_ast::ast::Identifier::Local( - when_str.clone(), - span.clone(), + // Update the when condition if provided + if let Some(when_str) = when { + watch_var.spec.when = match when_str.as_str() { + "manual" => crate::watch::WatchWhen::Manual, + "true" => crate::watch::WatchWhen::True, + _ => crate::watch::WatchWhen::FunctionName( + internal_baml_ast::ast::Identifier::Local( + when_str.clone(), + span.clone(), + ), ), - ), - }; - } + }; + } - watch_var.spec.span = span.clone(); - break; + watch_var.spec.span = span.clone(); + break; + } } } } + Statement::WatchNotify { variable, .. } => { + // Manually trigger a watch notification for this variable + fire_watch_notification_for_variable( + scopes, + variable, + watch_handler, + function_name, + )?; + } } - Statement::WatchNotify { variable, .. } => { - // Manually trigger a watch notification for this variable - fire_watch_notification_for_variable( - scopes, - variable, - watch_handler, - function_name, - )?; - } + Ok(None) + }) + } + + for stmt in block.statements.iter().take(statements_to_execute) { + let result = handle_statement( + stmt, + scopes, + thir, + run_llm_function, + watch_handler, + function_name, + ) + .await?; + + if let Some(control_flow) = result { + return Ok(control_flow); } } diff --git a/engine/baml-compiler/src/thir/typecheck.rs b/engine/baml-compiler/src/thir/typecheck.rs index 20aef6b8b4..e9e23de945 100644 --- a/engine/baml-compiler/src/thir/typecheck.rs +++ b/engine/baml-compiler/src/thir/typecheck.rs @@ -804,6 +804,21 @@ fn typecheck_statement( diagnostics: &mut Diagnostics, ) -> Option> { match stmt { + hir::Statement::AnnotatedStatement { headers, statement } => { + if let Some(statement) = statement { + typecheck_statement(statement, context, diagnostics).map(|stmt| { + thir::Statement::AnnotatedStatement { + headers: headers.clone(), + statement: Some(Box::new(stmt)), + } + }) + } else { + Some(thir::Statement::AnnotatedStatement { + headers: headers.clone(), + statement: None, + }) + } + } hir::Statement::Let { name, value, diff --git a/engine/baml-compiler/src/watch.rs b/engine/baml-compiler/src/watch.rs index 008ceddce0..bd6e9b0883 100644 --- a/engine/baml-compiler/src/watch.rs +++ b/engine/baml-compiler/src/watch.rs @@ -248,6 +248,14 @@ impl FunctionMetadata { diagnostics: &mut Diagnostics, ) { match statement { + thir::Statement::AnnotatedStatement { + statement, + headers: _, + } => { + if let Some(statement) = statement { + self.analyze_statement(statement, next_statement, diagnostics); + } + } thir::Statement::Let { value, watch, name, .. } => { diff --git a/engine/baml-compiler/src/watch/watch_event.rs b/engine/baml-compiler/src/watch/watch_event.rs index 5386a12dab..57772a3ca1 100644 --- a/engine/baml-compiler/src/watch/watch_event.rs +++ b/engine/baml-compiler/src/watch/watch_event.rs @@ -57,7 +57,11 @@ impl fmt::Display for WatchNotification { } }, WatchBamlValue::Block(label) => { - write!(f, "(block) {label}") + write!( + f, + "(block) {function_name}.{label}", + function_name = self.function_name + ) } WatchBamlValue::StreamStart(stream_id) => { write!(f, "(stream start) {stream_id}") diff --git a/engine/baml-runtime/src/async_vm_runtime.rs b/engine/baml-runtime/src/async_vm_runtime.rs index e79cbe8ead..c471d95d24 100644 --- a/engine/baml-runtime/src/async_vm_runtime.rs +++ b/engine/baml-runtime/src/async_vm_runtime.rs @@ -345,44 +345,57 @@ impl BamlAsyncVmRuntime { } } - Ok(VmExecState::Notify(nodes)) => { - for node in nodes { - let state = vm.watch.root_state(node).unwrap(); - let baml_vm::watch::NodeId::LocalVar(stack_index) = node else { - break 'mainloop Err(anyhow!("expected local variable notification, got object notification {:?}", node)); - }; - let (watched_var_name, function_name) = - vm.watched_vars.get(&stack_index).unwrap(); - baml_log::debug!("[VM] Notify: {}", &state.channel); - - let fake_meta = watch::WatchValueMetadata { - constraints: Vec::new(), - response_checks: Vec::new(), - completion: Completion::default(), - r#type: baml_types::TypeIR::Top(Default::default()), - }; - - let current_value = - try_baml_value_from_vm_value(&vm, &state.value).unwrap(); - - let baml_value_with_meta = - BamlValueWithMeta::with_const_meta(¤t_value, fake_meta); - - let notification = watch::WatchNotification::new_var( - watched_var_name.to_owned(), // variable name - state.channel.to_owned(), // channel name - baml_value_with_meta, - function_name.to_owned(), - ); + Ok(VmExecState::Notify(notification)) => { + match notification { + baml_vm::vm::WatchNotification::Variables(nodes) => { + for node in nodes { + let state = vm.watch.root_state(node).unwrap(); + let baml_vm::watch::NodeId::LocalVar(stack_index) = node else { + break 'mainloop Err(anyhow!("expected local variable notification, got object notification {:?}", node)); + }; + let (watched_var_name, function_name) = + vm.watched_vars.get(&stack_index).unwrap(); + baml_log::debug!("[VM] Notify: {}", &state.channel); + + let fake_meta = watch::WatchValueMetadata { + constraints: Vec::new(), + response_checks: Vec::new(), + completion: Completion::default(), + r#type: baml_types::TypeIR::Top(Default::default()), + }; + + let current_value = + try_baml_value_from_vm_value(&vm, &state.value).unwrap(); + + let baml_value_with_meta = + BamlValueWithMeta::with_const_meta(¤t_value, fake_meta); + + let notification = watch::WatchNotification::new_var( + watched_var_name.to_owned(), // variable name + state.channel.to_owned(), // channel name + baml_value_with_meta, + function_name.to_owned(), + ); - if let Some(handler) = watch_handler.as_mut() { - handler(notification); + if let Some(handler) = watch_handler.as_mut() { + handler(notification); + } + } + } + baml_vm::vm::WatchNotification::Block(notification) => { + if let Some(handler) = watch_handler.as_mut() { + handler(watch::WatchNotification::new_block( + String::from_utf8_lossy(¬ification.block_name).to_string(), + String::from_utf8_lossy(¬ification.function_name) + .to_string(), + )); + } } } } - Ok(VmExecState::ScheduleFuture(idx)) => { - let pending_future = match vm.pending_future(idx) { + Ok(VmExecState::ScheduleFuture(idxan)) => { + let pending_future = match vm.pending_future(idxan) { Ok(f) => f, Err(e) => { break 'mainloop Err( @@ -459,7 +472,7 @@ impl BamlAsyncVmRuntime { .await; // TODO: Handle panic somehow. - futures_tx.send((idx, result)).unwrap_or_else(|e| { + futures_tx.send((idxan, result)).unwrap_or_else(|e| { panic!("failed to send LLM function result to futures channel: {e}") }); } @@ -625,7 +638,7 @@ impl BamlAsyncVmRuntime { ); // TODO: Handle panic somehow. - futures_tx.send((idx, (Ok(result), current_call_id))).unwrap_or_else(|e| { + futures_tx.send((idxan, (Ok(result), current_call_id))).unwrap_or_else(|e| { panic!("failed to send LLM function result to futures channel: {e}") }); } diff --git a/engine/baml-vm/src/bytecode.rs b/engine/baml-vm/src/bytecode.rs index 818587d7f2..9c47743806 100644 --- a/engine/baml-vm/src/bytecode.rs +++ b/engine/baml-vm/src/bytecode.rs @@ -230,6 +230,30 @@ pub enum Instruction { /// /// Format: `ASSERT` Assert, + + /// Notifies about entering or exiting a block. + /// + /// Format: `NOTIFY_BLOCK function_name block_name` + NotifyBlock(BlockNotification), +} + +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct BlockNotification { + // This is a hack cause i don't wanna deal with implementing Copy by allocating this on the heap + doing an object pointer. + pub function_name: [u8; 1024], + pub block_name: [u8; 1024], + pub level: usize, + pub block_type: BlockNotificationType, + pub is_enter: bool, +} + +#[derive(Clone, Copy, Debug, PartialEq)] +pub enum BlockNotificationType { + Statement, + If, + While, + For, + Function, } #[derive(Clone, Copy, Debug, PartialEq)] @@ -335,6 +359,19 @@ impl std::fmt::Display for Instruction { Instruction::Assert => f.write_str("ASSERT"), Instruction::AllocMap(n) => write!(f, "ALLOC_MAP {n}"), Instruction::Watch => f.write_str("WATCH"), + Instruction::NotifyBlock(notification) => { + write!( + f, + "{}_BLOCK {function_name}.{block_name}", + if notification.is_enter { + "ENTER" + } else { + "EXIT" + }, + function_name = String::from_utf8_lossy(¬ification.function_name), + block_name = String::from_utf8_lossy(¬ification.block_name), + ) + } } } } diff --git a/engine/baml-vm/src/debug.rs b/engine/baml-vm/src/debug.rs index fe312cd462..91193e4005 100644 --- a/engine/baml-vm/src/debug.rs +++ b/engine/baml-vm/src/debug.rs @@ -57,6 +57,9 @@ pub fn display_instruction( let instruction = &function.bytecode.instructions[instruction_ptr as usize]; let metadata = match instruction { + Instruction::NotifyBlock(notification) => { + format!("({})", String::from_utf8_lossy(¬ification.block_name)) + } Instruction::LoadConst(index) => format!( "({})", display_value(&function.bytecode.constants[*index], objects) @@ -179,6 +182,7 @@ const COLUMN_MARGIN: usize = 3; /// Get color for instruction based on its type fn instruction_color(instruction: &Instruction) -> Color { match instruction { + Instruction::NotifyBlock(_) => Color::BrightYellow, Instruction::LoadConst(_) | Instruction::LoadVar(_) | Instruction::LoadGlobal(_) diff --git a/engine/baml-vm/src/vm.rs b/engine/baml-vm/src/vm.rs index 955feec6c0..e6bad28b67 100644 --- a/engine/baml-vm/src/vm.rs +++ b/engine/baml-vm/src/vm.rs @@ -3,7 +3,7 @@ use std::collections::HashMap; use baml_types::{BamlMap, BamlMedia}; use crate::{ - bytecode::{BinOp, CmpOp, Instruction}, + bytecode::{BinOp, BlockNotification, CmpOp, Instruction}, errors::{ErrorLocation, InternalError, RuntimeError, VmError}, indexable::{EvalStack, GlobalPool, ObjectIndex, ObjectPool, StackIndex}, types::{ @@ -220,7 +220,13 @@ pub enum VmExecState { Complete(Value), /// Notify about watched variables. - Notify(Vec), + Notify(WatchNotification), +} + +#[derive(Debug, PartialEq)] +pub enum WatchNotification { + Variables(Vec), + Block(BlockNotification), } #[derive(Clone, Debug)] @@ -470,6 +476,9 @@ impl Vm { // } match function.bytecode.instructions[instruction_ptr as usize] { + Instruction::NotifyBlock(notification) => { + return Ok(VmExecState::Notify(WatchNotification::Block(notification))); + } Instruction::LoadConst(index) => { let value = &function.bytecode.constants[index]; self.stack.push(*value); @@ -527,7 +536,9 @@ impl Vm { (NodeId::HeapObject(a), NodeId::HeapObject(b)) => a.cmp(b), }); if !notifications.is_empty() { - return Ok(VmExecState::Notify(notifications)); + return Ok(VmExecState::Notify(WatchNotification::Variables( + notifications, + ))); } } } @@ -616,7 +627,9 @@ impl Vm { }); if !notifications.is_empty() { - return Ok(VmExecState::Notify(notifications)); + return Ok(VmExecState::Notify(WatchNotification::Variables( + notifications, + ))); } } @@ -1111,7 +1124,9 @@ impl Vm { (NodeId::HeapObject(a), NodeId::HeapObject(b)) => a.cmp(b), }); if !notifications.is_empty() { - return Ok(VmExecState::Notify(notifications)); + return Ok(VmExecState::Notify(WatchNotification::Variables( + notifications, + ))); } } @@ -1182,7 +1197,9 @@ impl Vm { (NodeId::HeapObject(a), NodeId::HeapObject(b)) => a.cmp(b), }); if !notifications.is_empty() { - return Ok(VmExecState::Notify(notifications)); + return Ok(VmExecState::Notify(WatchNotification::Variables( + notifications, + ))); } } diff --git a/engine/baml-vm/tests/common.rs b/engine/baml-vm/tests/common.rs index c5f3d244d8..ff7e0cb275 100644 --- a/engine/baml-vm/tests/common.rs +++ b/engine/baml-vm/tests/common.rs @@ -183,10 +183,13 @@ impl ExecState { Value::from_vm_value(&value, vm).map(ExecState::Complete) } VmExecState::Notify(roots) => { - let nodes = roots - .iter() - .map(|node_id| Notification::from_node_id(node_id, vm)) - .collect::>>()?; + let nodes = match roots { + baml_vm::vm::WatchNotification::Variables(nodes) => nodes + .iter() + .map(|node_id| Notification::from_node_id(node_id, vm)) + .collect::>>()?, + baml_vm::vm::WatchNotification::Block(_) => todo!("block notifications"), + }; Ok(ExecState::Emit(nodes)) } } diff --git a/integ-tests/baml_src/test-files/workflows/workflow_emit_simple.baml b/integ-tests/baml_src/test-files/workflows/workflow_emit_simple.baml index 9e7ba2b869..9e16217be4 100644 --- a/integ-tests/baml_src/test-files/workflows/workflow_emit_simple.baml +++ b/integ-tests/baml_src/test-files/workflows/workflow_emit_simple.baml @@ -5,19 +5,35 @@ function NotEmpty(value: string) -> bool { // Simple test without loops function SimpleWatchWithFilter() -> int { + //# Section1 watch let word: string = ""; word.$watch.options(baml.WatchOptions{ when: NotEmpty }); + //## Section2 word = "hello"; // Should notify (not empty) + //# Section3 word = "world"; // Should notify (not empty) word = ""; // Should NOT notify (empty) word = "test"; // Should notify (not empty) + + //# Section4 word.$watch.options(baml.WatchOptions{ channel: "new_name" }); word = "with_new_name"; // Should notify (not empty) word.$watch.notify(); - 42 + //# SectionIf + let x = if (false) { + //# Section42 + 42 + } else if (true) { + //# Section41 + 41 + } else { + 40 + }; + + x } test TestSimpleWatchWithFilter() {