baml-lanugage/fix: use build_plan for llm functions, delete execute

Author: sxlijin

baml_language/crates/baml_builtins2/baml_std/baml/llm.baml

@@ -3,11 +3,6 @@ class OrchestrationStep {
   delay_ms int
 }
 
-class ExecutionResult {
-  ok bool
-  value unknown
-}
-
 class ExecutionContext {
   jinja_string string
   args map<string, unknown>
@@ -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();
+
+  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" };
 }

baml_language/crates/baml_builtins2/baml_std/baml/llm_types.baml

@@ -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;
+      },
+    }
+  }
+
   function get_constructor(
     self
   ) -> () -> PrimitiveClient throws root.errors.InvalidArgument {
@@ -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 {
@@ -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
 }

baml_language/crates/baml_compiler2_tir/src/builder.rs

@@ -2235,6 +2235,32 @@ impl<'db> TypeInferenceBuilder<'db> {
                     // reported here — they'll be reported at the definition site.
                     ty
                 }
+                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,
+                    }
+                }
                 _ => Ty::Unknown,
             }
         } else if let Some(def) = self.package_items.lookup_type(&lookup_path) {