Skip to content

baml-language/fix: use build_plan for llm functions, delete execute #3279

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
sxlijin wants to merge 1 commit into
base: canary
Choose a base branch
from
+139 −237
Open

baml-language/fix: use build_plan for llm functions, delete execute #3279

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 15 additions & 7 deletions baml_language/crates/baml_builtins2/baml_std/baml/llm.baml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,11 +3,6 @@ class OrchestrationStep {
delay_ms int
}

class ExecutionResult {
ok bool
value unknown
}

class ExecutionContext {
jinja_string string
args map<string, unknown>
Expand Down Expand Up @@ -58,6 +53,19 @@ function call_llm_function<T>(client: Client, function_name: string, args: map<s
function_name: function_name,
};

let result: T = client.execute(context, 0);
result
let steps = client.build_plan();
client.advance_round_robin();
Copy link
Copy Markdown

@chatgpt-codex-connector chatgpt-codex-connector bot Mar 25, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Advance round-robin counters only for attempted steps

call_llm_function now mutates round-robin state before any step runs, which breaks fallback behavior when an earlier branch succeeds. Because build_plan() flattens all fallback branches but the loop returns on the first successful step, many planned steps are never attempted; however client.advance_round_robin() still increments counters in those untouched subtrees, so future calls can skip providers that were never actually used. This is a regression from the previous execute-on-visit behavior and changes routing deterministically for Fallback[..., RoundRobin[...]] clients.

Useful? React with 👍 / 👎.


for (let step in steps) {
if (step.delay_ms > 0) {
root.sys.sleep(step.delay_ms);
}

let result: T = execute_step(step, context) catch (e) {
_ => { continue; }
};
return result;
}

throw root.errors.DevOther { message: "All orchestration steps failed" };
Comment on lines +64 to +70
Copy link
Copy Markdown
Contributor

@coderabbitai coderabbitai bot Mar 25, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

Don't treat every execute_step() failure as retryable.

execute_step() can fail with deterministic configuration errors too, not just provider failures. Catching _ here retries the next step and eventually replaces the real cause with "All orchestration steps failed", which can also send extra requests for a non-retryable bug. Restrict the fallback path to step-local provider/HTTP failures and rethrow the rest.

}
135 changes: 32 additions & 103 deletions baml_language/crates/baml_builtins2/baml_std/baml/llm_types.baml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,20 @@ class Client {
}
}

function advance_round_robin(self) -> void {
match (self.client_type) {
ClientType.Primitive => {},
ClientType.Fallback => {
for (let sub in self.sub_clients) {
sub.advance_round_robin();
}
},
ClientType.RoundRobin => {
self.counter += 1;
},
}
}
Comment on lines +35 to +47
Copy link
Copy Markdown
Contributor

@coderabbitai coderabbitai bot Mar 25, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

Advance the selected round-robin child too.

build_plan_with_state() consumes the chosen child subtree after computing idx, so nested round-robin nodes under that child can participate in the plan. advance_round_robin() only increments self.counter, though, so those nested counters never advance and the same inner branch is reused whenever this outer RR picks that child again. Mirror the same child selection here and recurse into the chosen sub-client as well.


function get_constructor(
self
) -> () -> PrimitiveClient throws root.errors.InvalidArgument {
Expand Down Expand Up @@ -120,109 +134,6 @@ class Client {
}
}

function execute<T>(
self,
context: ExecutionContext,
inherited_delay_ms: int,
) -> T {
match (self.retry) {
r: RetryPolicy => {
let current_delay = r.initial_delay_ms + 0.0

for (let attempt = 0; attempt <= r.max_retries; attempt += 1) {
let attempt_delay = inherited_delay_ms
if (attempt > 0) {
attempt_delay = root.math.trunc(current_delay)
let next = current_delay * r.multiplier
if (next > r.max_delay_ms + 0.0) {
current_delay = r.max_delay_ms + 0.0
} else {
current_delay = next
}
}

if (attempt == r.max_retries) {
attempt_delay = inherited_delay_ms
}

let result2: T = self.execute_once(
context,
attempt_delay,
) catch (e) {
_ => { continue; }
};

return result2;
}

throw root.errors.DevOther { message: "All orchestration steps failed" };
}
null => {
let result: T = self.execute_once(
context,
inherited_delay_ms,
);
return result;
},
}
}

function execute_once<T>(
self,
context: ExecutionContext,
active_delay_ms: int,
) -> T {
match (self.client_type) {
ClientType.Primitive => {
let resolve_fn = self.get_constructor()
let primitive = resolve_fn()

let prompt = primitive.render_prompt(context.jinja_string, context.args)
let specialized = primitive.specialize_prompt(prompt)
let http_request = primitive.build_request(specialized)
let http_response = root.http.send(http_request)

if (http_response.ok()) {
let body = http_response.text()
let return_type = get_return_type(context.function_name)
let result: T = primitive.parse(body, return_type) catch (e) {
_ => {
if (active_delay_ms > 0) {
root.sys.sleep(active_delay_ms)
}
throw e;
}
};
result
} else {
throw root.errors.DevOther { message: "HTTP request failed" };
}
}

ClientType.Fallback => {
for (let sub in self.sub_clients) {
let result2: T = sub.execute(
context,
active_delay_ms,
) catch (e) {
_ => { continue; }
};
return result2;
}
throw root.errors.DevOther { message: "All orchestration steps failed" };
}

ClientType.RoundRobin => {
let idx = self.counter % self.sub_clients.length()
self.counter += 1
let result3: T = self.sub_clients[idx].execute(
context,
active_delay_ms,
);
return result3;
}
}
}
}

class PrimitiveClientOptions {
Expand Down Expand Up @@ -285,6 +196,24 @@ class PrimitiveClient {
}
}

function execute_step<T>(
step: OrchestrationStep,
context: ExecutionContext,
) -> T {
let prompt = step.primitive_client.render_prompt(context.jinja_string, context.args);
let specialized = step.primitive_client.specialize_prompt(prompt);
let http_request = step.primitive_client.build_request(specialized);
let http_response = root.http.send(http_request);

if (http_response.ok()) {
let body = http_response.text();
let return_type = get_return_type(context.function_name);
step.primitive_client.parse(body, return_type)
} else {
throw root.errors.DevOther { message: "HTTP request failed" };
}
}

function get_jinja_template(function_name: string) -> string throws root.errors.InvalidArgument {
$rust_io_function
}
Expand Down
26 changes: 26 additions & 0 deletions baml_language/crates/baml_compiler2_tir/src/builder.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2258,6 +2258,32 @@ impl<'db> TypeInferenceBuilder<'db> {
),
}
}
Definition::Let(let_loc) => {
// Determine type from the let-binding's origin.
let db = self.context.db();
let item_tree =
baml_compiler2_hir::file_item_tree(db, let_loc.file(db));
let let_data = &item_tree[let_loc.id(db)];
match let_data.origin {
baml_compiler2_ast::ast::LetOrigin::Client => {
// client<llm> declarations produce Client instances.
Ty::Class(crate::ty::QualifiedTypeName::new(
baml_base::Name::new("baml"),
vec![baml_base::Name::new("llm")],
baml_base::Name::new("Client"),
))
}
baml_compiler2_ast::ast::LetOrigin::RetryPolicy => {
// retry_policy declarations produce RetryPolicy instances.
Ty::Class(crate::ty::QualifiedTypeName::new(
baml_base::Name::new("baml"),
vec![baml_base::Name::new("llm")],
baml_base::Name::new("RetryPolicy"),
))
}
_ => Ty::Unknown,
}
Comment on lines +2261 to +2285
Copy link
Copy Markdown
Contributor

@coderabbitai coderabbitai bot Mar 25, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

Handle let-bound globals in package-qualified paths too.

This branch only fixes bare identifiers through infer_single_name(). root.A.build_plan() and package-qualified client/retry-policy references still go through resolve_package_item(), which only recognizes Definition::Function, so they continue to infer as Ty::Unknown and member lookup fails. Please share the same Definition::Let → type mapping with the multi-segment/package path resolver.

}
_ => Ty::Unknown,
}
} else if let Some(def) = self.package_items.lookup_type(&lookup_path) {
Expand Down
Loading
Loading