Skip to content

Commit c322b64

Browse files
imalsogregimalsogreg
authored
remove the BamlRuntime/InternalBamlRuntime split (BoundaryML#2588)
The API boundary originally established between `BamlRuntime` and `InternalBamlRuntime` hasn't been followed, and so the internal runtime is now used in many non-internal places. The split between these two APIs became confusing. So this PR just removes `InternalBamlRuntime` and make `BamlRuntime` the main entrypoint to the runtime. In the future, we may want to rename this to `LLMRuntime`, to highlight the fact that it only serves LLM functions, not more general expression functions. Expression functions need their own runtime, which is often aliased to `CoreBamlRuntime`. Renaming `BamlRuntime` to `LLMRuntime` would remove this additional confusing name.
1 parent e4ae745 commit c322b64

File tree

28 files changed

+754
-382
lines changed

28 files changed

+754
-382
lines changed

‎engine/baml-runtime/src/async_interpreter_runtime.rs‎

Lines changed: 10 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,6 @@ use crate::on_log_event::LogEventCallbackSync;
2222
use crate::{
2323
client_registry::ClientRegistry,
2424
internal::llm_client::{orchestrator::OrchestrationScope, LLMResponse},
25-
runtime::InternalBamlRuntime,
2625
runtime_interface::ExperimentalTracingInterface,
2726
tracing::TracingCall,
2827
tracingv2::storage::storage::Collector,
@@ -59,7 +58,7 @@ impl TryFrom<LlmRuntime> for BamlAsyncInterpreterRuntime {
5958
let async_runtime = Arc::clone(&llm_runtime.async_runtime);
6059

6160
// Stage 1: AST -> HIR
62-
let hir_program = hir::Hir::from_ast(&llm_runtime.inner.db.ast);
61+
let hir_program = hir::Hir::from_ast(&llm_runtime.db.ast);
6362

6463
// Stage 2: HIR -> THIR (typecheck)
6564
let mut diagnostics = Diagnostics::new("dummy".into());
@@ -81,8 +80,8 @@ impl TryFrom<LlmRuntime> for BamlAsyncInterpreterRuntime {
8180
}
8281

8382
impl BamlAsyncInterpreterRuntime {
84-
pub fn internal(&self) -> &Arc<InternalBamlRuntime> {
85-
&self.llm_runtime.inner
83+
pub fn internal(&self) -> &LlmRuntime {
84+
&self.llm_runtime
8685
}
8786

8887
pub fn disassemble(&self, function_name: &str) {
@@ -236,7 +235,6 @@ impl BamlAsyncInterpreterRuntime {
236235
async move {
237236
// Find the LLM function to get parameter names
238237
let llm_fn = llm_runtime
239-
.inner
240238
.ir()
241239
.find_function(&fn_name)
242240
.map_err(|e| anyhow::anyhow!("LLM function not found: {}: {}", fn_name, e))?;
@@ -596,7 +594,6 @@ impl BamlAsyncInterpreterRuntime {
596594
test_name: &str,
597595
) -> Result<Option<TypeBuilder>> {
598596
self.llm_runtime
599-
.inner
600597
.get_test_type_builder(function_name, test_name)
601598
}
602599

@@ -759,42 +756,42 @@ fn baml_value_with_meta_to_baml_value(
759756

760757
impl crate::runtime_interface::InternalRuntimeInterface for BamlAsyncInterpreterRuntime {
761758
fn features(&self) -> crate::internal::ir_features::IrFeatures {
762-
self.llm_runtime.inner.features()
759+
self.llm_runtime.features()
763760
}
764761

765762
fn diagnostics(&self) -> &internal_baml_core::internal_baml_diagnostics::Diagnostics {
766-
self.llm_runtime.inner.diagnostics()
763+
self.llm_runtime.diagnostics()
767764
}
768765

769766
fn orchestration_graph(
770767
&self,
771768
client_name: &internal_llm_client::ClientSpec,
772769
ctx: &crate::runtime_context::RuntimeContext,
773770
) -> Result<Vec<crate::internal::llm_client::orchestrator::OrchestratorNode>> {
774-
self.llm_runtime.inner.orchestration_graph(client_name, ctx)
771+
self.llm_runtime.orchestration_graph(client_name, ctx)
775772
}
776773

777774
fn function_graph(
778775
&self,
779776
function_name: &str,
780777
ctx: &crate::runtime_context::RuntimeContext,
781778
) -> Result<String> {
782-
self.llm_runtime.inner.function_graph(function_name, ctx)
779+
self.llm_runtime.function_graph(function_name, ctx)
783780
}
784781

785782
fn get_function<'ir>(
786783
&'ir self,
787784
function_name: &str,
788785
) -> Result<internal_baml_core::ir::FunctionWalker<'ir>> {
789-
self.llm_runtime.inner.get_function(function_name)
786+
self.llm_runtime.get_function(function_name)
790787
}
791788

792789
fn get_expr_function<'ir>(
793790
&'ir self,
794791
function_name: &str,
795792
ctx: &crate::runtime_context::RuntimeContext,
796793
) -> Result<internal_baml_core::ir::ExprFunctionWalker<'ir>> {
797-
self.llm_runtime.inner.get_expr_function(function_name, ctx)
794+
self.llm_runtime.get_expr_function(function_name, ctx)
798795
}
799796

800797
async fn render_prompt(
@@ -809,7 +806,6 @@ impl crate::runtime_interface::InternalRuntimeInterface for BamlAsyncInterpreter
809806
internal_llm_client::AllowedRoleMetadata,
810807
)> {
811808
self.llm_runtime
812-
.inner
813809
.render_prompt(function_name, ctx, params, node_index)
814810
.await
815811
}
@@ -823,13 +819,12 @@ impl crate::runtime_interface::InternalRuntimeInterface for BamlAsyncInterpreter
823819
node_index: Option<usize>,
824820
) -> Result<String> {
825821
self.llm_runtime
826-
.inner
827822
.render_raw_curl(function_name, ctx, prompt, render_settings, node_index)
828823
.await
829824
}
830825

831826
fn ir(&self) -> &internal_baml_core::ir::repr::IntermediateRepr {
832-
self.llm_runtime.inner.ir()
827+
self.llm_runtime.ir()
833828
}
834829

835830
fn get_test_params(
@@ -840,7 +835,6 @@ impl crate::runtime_interface::InternalRuntimeInterface for BamlAsyncInterpreter
840835
strict: bool,
841836
) -> Result<BamlMap<String, BamlValue>> {
842837
self.llm_runtime
843-
.inner
844838
.get_test_params(function_name, test_name, ctx, strict)
845839
}
846840

@@ -851,7 +845,6 @@ impl crate::runtime_interface::InternalRuntimeInterface for BamlAsyncInterpreter
851845
ctx: &crate::runtime_context::RuntimeContext,
852846
) -> Result<Vec<baml_types::Constraint>> {
853847
self.llm_runtime
854-
.inner
855848
.get_test_constraints(function_name, test_name, ctx)
856849
}
857850

@@ -861,7 +854,6 @@ impl crate::runtime_interface::InternalRuntimeInterface for BamlAsyncInterpreter
861854
test_name: &str,
862855
) -> Result<Option<TypeBuilder>> {
863856
self.llm_runtime
864-
.inner
865857
.get_test_type_builder(function_name, test_name)
866858
}
867859
}

‎engine/baml-runtime/src/async_vm_runtime.rs‎

Lines changed: 12 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,6 @@ use crate::on_log_event::LogEventCallbackSync;
2525
use crate::{
2626
client_registry::ClientRegistry,
2727
internal::llm_client::{orchestrator::OrchestrationScope, LLMResponse},
28-
runtime::InternalBamlRuntime,
2928
runtime_interface::ExperimentalTracingInterface,
3029
tracing::TracingCall,
3130
tracingv2::storage::storage::Collector,
@@ -65,7 +64,7 @@ impl TryFrom<LlmRuntime> for BamlAsyncVmRuntime {
6564
#[cfg(not(target_arch = "wasm32"))]
6665
let async_runtime = Arc::clone(&llm_runtime.async_runtime);
6766

68-
let program = baml_compiler::compile(&llm_runtime.inner.db)?;
67+
let program = baml_compiler::compile(&llm_runtime.db)?;
6968

7069
Ok(Self {
7170
llm_runtime: Arc::new(llm_runtime),
@@ -78,8 +77,8 @@ impl TryFrom<LlmRuntime> for BamlAsyncVmRuntime {
7877
}
7978

8079
impl BamlAsyncVmRuntime {
81-
pub fn internal(&self) -> &Arc<InternalBamlRuntime> {
82-
&self.llm_runtime.inner
80+
pub fn internal(&self) -> &LlmRuntime {
81+
&self.llm_runtime
8382
}
8483

8584
pub fn disassemble(&self, function_name: &str) {
@@ -212,7 +211,6 @@ impl BamlAsyncVmRuntime {
212211

213212
let Some(expr_fn) = self
214213
.llm_runtime
215-
.inner
216214
.ir()
217215
.expr_fns
218216
.iter()
@@ -358,7 +356,6 @@ impl BamlAsyncVmRuntime {
358356
baml_vm::FutureKind::Llm => {
359357
let llm_fn = match self
360358
.llm_runtime
361-
.inner
362359
.ir()
363360
.find_function(&pending_future.function)
364361
{
@@ -519,7 +516,7 @@ impl BamlAsyncVmRuntime {
519516
let current_call_id = current_call_id.to_owned();
520517

521518
let output_format = jsonish::helpers::render_output_format(
522-
&self.llm_runtime.inner.ir,
519+
&self.llm_runtime.ir,
523520
&parse_as_type,
524521
&baml_types::EvaluationContext::default(),
525522
baml_types::StreamingMode::NonStreaming,
@@ -915,8 +912,7 @@ impl BamlAsyncVmRuntime {
915912
test_name: &str,
916913
) -> Result<Option<TypeBuilder>, anyhow::Error> {
917914
self.llm_runtime
918-
.inner
919-
.get_test_type_builder(function_name, test_name)
915+
.get_test_type_builder_impl(function_name, test_name)
920916
}
921917

922918
pub fn tracer_wrapper(&self) -> &Arc<BamlTracerWrapper> {
@@ -1206,42 +1202,42 @@ fn try_vm_value_from_function_result(
12061202

12071203
impl crate::runtime_interface::InternalRuntimeInterface for BamlAsyncVmRuntime {
12081204
fn features(&self) -> crate::internal::ir_features::IrFeatures {
1209-
self.llm_runtime.inner.features()
1205+
self.llm_runtime.features()
12101206
}
12111207

12121208
fn diagnostics(&self) -> &internal_baml_core::internal_baml_diagnostics::Diagnostics {
1213-
self.llm_runtime.inner.diagnostics()
1209+
self.llm_runtime.diagnostics()
12141210
}
12151211

12161212
fn orchestration_graph(
12171213
&self,
12181214
client_name: &internal_llm_client::ClientSpec,
12191215
ctx: &crate::runtime_context::RuntimeContext,
12201216
) -> anyhow::Result<Vec<crate::internal::llm_client::orchestrator::OrchestratorNode>> {
1221-
self.llm_runtime.inner.orchestration_graph(client_name, ctx)
1217+
self.llm_runtime.orchestration_graph(client_name, ctx)
12221218
}
12231219

12241220
fn function_graph(
12251221
&self,
12261222
function_name: &str,
12271223
ctx: &crate::runtime_context::RuntimeContext,
12281224
) -> anyhow::Result<String> {
1229-
self.llm_runtime.inner.function_graph(function_name, ctx)
1225+
self.llm_runtime.function_graph(function_name, ctx)
12301226
}
12311227

12321228
fn get_function<'ir>(
12331229
&'ir self,
12341230
function_name: &str,
12351231
) -> anyhow::Result<internal_baml_core::ir::FunctionWalker<'ir>> {
1236-
self.llm_runtime.inner.get_function(function_name)
1232+
self.llm_runtime.get_function(function_name)
12371233
}
12381234

12391235
fn get_expr_function<'ir>(
12401236
&'ir self,
12411237
function_name: &str,
12421238
ctx: &crate::runtime_context::RuntimeContext,
12431239
) -> anyhow::Result<internal_baml_core::ir::ExprFunctionWalker<'ir>> {
1244-
self.llm_runtime.inner.get_expr_function(function_name, ctx)
1240+
self.llm_runtime.get_expr_function(function_name, ctx)
12451241
}
12461242

12471243
async fn render_prompt(
@@ -1256,7 +1252,6 @@ impl crate::runtime_interface::InternalRuntimeInterface for BamlAsyncVmRuntime {
12561252
internal_llm_client::AllowedRoleMetadata,
12571253
)> {
12581254
self.llm_runtime
1259-
.inner
12601255
.render_prompt(function_name, ctx, params, node_index)
12611256
.await
12621257
}
@@ -1270,13 +1265,12 @@ impl crate::runtime_interface::InternalRuntimeInterface for BamlAsyncVmRuntime {
12701265
node_index: Option<usize>,
12711266
) -> anyhow::Result<String> {
12721267
self.llm_runtime
1273-
.inner
12741268
.render_raw_curl(function_name, ctx, prompt, render_settings, node_index)
12751269
.await
12761270
}
12771271

12781272
fn ir(&self) -> &internal_baml_core::ir::repr::IntermediateRepr {
1279-
self.llm_runtime.inner.ir()
1273+
self.llm_runtime.ir()
12801274
}
12811275

12821276
fn get_test_params(
@@ -1287,7 +1281,6 @@ impl crate::runtime_interface::InternalRuntimeInterface for BamlAsyncVmRuntime {
12871281
strict: bool,
12881282
) -> anyhow::Result<BamlMap<String, BamlValue>> {
12891283
self.llm_runtime
1290-
.inner
12911284
.get_test_params(function_name, test_name, ctx, strict)
12921285
}
12931286

@@ -1298,7 +1291,6 @@ impl crate::runtime_interface::InternalRuntimeInterface for BamlAsyncVmRuntime {
12981291
ctx: &crate::runtime_context::RuntimeContext,
12991292
) -> anyhow::Result<Vec<baml_types::Constraint>> {
13001293
self.llm_runtime
1301-
.inner
13021294
.get_test_constraints(function_name, test_name, ctx)
13031295
}
13041296

@@ -1308,7 +1300,6 @@ impl crate::runtime_interface::InternalRuntimeInterface for BamlAsyncVmRuntime {
13081300
test_name: &str,
13091301
) -> anyhow::Result<Option<TypeBuilder>> {
13101302
self.llm_runtime
1311-
.inner
13121303
.get_test_type_builder(function_name, test_name)
13131304
}
13141305
}

‎engine/baml-runtime/src/cli/check.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -64,7 +64,7 @@ impl CheckArgs {
6464
exit(1);
6565
}
6666
Ok(runtime) => {
67-
let diagnostics = runtime.inner.diagnostics();
67+
let diagnostics = &runtime.diagnostics;
6868
if diagnostics.has_warnings() {
6969
println!("{}", diagnostics.warnings_to_pretty_string());
7070
}

‎engine/baml-runtime/src/cli/generate.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,7 @@ impl GenerateArgs {
6060
.context("Failed to build BAML runtime")?;
6161

6262
// Display warnings only if the feature flag is enabled
63-
let diagnostics = runtime.inner.diagnostics();
63+
let diagnostics = &runtime.diagnostics;
6464
if feature_flags.should_display_warnings() && diagnostics.has_warnings() {
6565
eprintln!("{}", diagnostics.warnings_to_pretty_string());
6666
}

‎engine/baml-runtime/src/cli/repl.rs‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -272,7 +272,7 @@ impl ReplState {
272272
}
273273
let mut names = Vec::new();
274274
if let Some(runtime) = &self.runtime {
275-
let internal = &runtime.inner;
275+
let internal = &runtime;
276276
for function in internal.db.walk_functions() {
277277
names.push(function.name().to_string());
278278
}
@@ -285,7 +285,7 @@ impl ReplState {
285285
fn function_parameters(&self) -> Result<HashMap<String, Vec<String>>> {
286286
let hir = match self.runtime.as_ref() {
287287
Some(runtime) => {
288-
let internal = &runtime.inner;
288+
let internal = &runtime;
289289
Hir::from_ast(&internal.db.ast)
290290
}
291291
None => Hir::empty(),
@@ -311,7 +311,7 @@ impl ReplState {
311311
.as_ref()
312312
.ok_or_else(|| anyhow!("No BAML sources loaded. Use :load <path> to load sources."))?;
313313

314-
let internal = &runtime.inner;
314+
let internal = &runtime;
315315

316316
// Convert AST to HIR
317317
let hir = Hir::from_ast(&internal.db.ast);
@@ -460,7 +460,7 @@ impl ReplState {
460460
let hir = match self.runtime.as_ref() {
461461
Some(runtime) => {
462462
// Get the internal runtime to access the existing context
463-
let internal = &runtime.inner;
463+
let internal = &runtime;
464464

465465
// Convert AST to HIR from existing loaded sources
466466
Hir::from_ast(&internal.db.ast)
@@ -596,7 +596,7 @@ impl ReplState {
596596
let hir = match self.runtime.as_ref() {
597597
Some(runtime) => {
598598
// Get the internal runtime to access the existing context
599-
let internal = &runtime.inner;
599+
let internal = &runtime;
600600

601601
// Convert AST to HIR from existing loaded sources
602602
Hir::from_ast(&internal.db.ast)
@@ -679,7 +679,7 @@ impl ReplState {
679679

680680
// Add functions and declarations from THIR if available
681681
if let Some(runtime) = &self.runtime {
682-
let internal = &runtime.inner;
682+
let internal = &runtime;
683683

684684
// Add function names
685685
for function in internal.db.walk_functions() {

‎engine/baml-runtime/src/cli/serve/mod.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -640,7 +640,7 @@ Streaming is available via http://localhost:{port}/stream/{{FunctionName}}, but
640640
.map_err(|_| BamlError::InternalError {
641641
message: "Failed to make placeholder generator".to_string(),
642642
})?;
643-
let schema: OpenApiSchema = OpenApiSchema::from_ir(locked.inner.ir.as_ref());
643+
let schema: OpenApiSchema = OpenApiSchema::from_ir(locked.ir.as_ref());
644644
serde_json::to_string(&schema).map_err(|e| {
645645
log::warn!("Failed to serialize openapi schema: {e}");
646646
BamlError::InternalError {

0 commit comments

Comments
 (0)