Skip to content

Commit 4397cf1

Browse files
ba11b0yaaronvgimalsogreg
authored
lazy load env vars (#1949)
- allows lazy loading for env vars on function calls - removes env vars from ctx manager - supporting changes for cli, playground and codegen for all languages. - refactors tracer to use dashmap and reuse until required env vars change TODO: @aaronvg to check AWS SSL on playground and ruby --------- Co-authored-by: aaronvg <aaron@boundaryml.com> Co-authored-by: Greg Hale <ImAlsoGreg@gmail.com>
1 parent 9212af2 commit 4397cf1

102 files changed

Lines changed: 26494 additions & 6090 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎.github/workflows/primary.yml‎

Lines changed: 16 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -81,25 +81,33 @@ jobs:
8181
run: pnpm install --frozen-lockfile
8282
working-directory: integ-tests/typescript
8383
- name: Test Node Generator
84-
run: pnpm generate
84+
run: |
85+
pnpm generate || (echo "::error:: merge canary and run codegen again" && exit 1)
8586
working-directory: integ-tests/typescript
8687
- name: Ensure No Changes
87-
run: git diff --exit-code
88+
run: |
89+
git diff --exit-code || (echo "::error:: merge canary and run codegen again" && exit 1)
8890
- name: Test Node Generator (1 of 3)
89-
run: pnpm generate
91+
run: |
92+
pnpm generate || (echo "::error:: merge canary and run codegen again" && exit 1)
9093
working-directory: integ-tests/typescript
9194
- name: Ensure No Changes (1 of 3)
92-
run: git diff --exit-code
95+
run: |
96+
git diff --exit-code || (echo "::error:: merge canary and run codegen again" && exit 1)
9397
- name: Test Node Generator (2 of 3)
94-
run: pnpm generate
98+
run: |
99+
pnpm generate || (echo "::error:: merge canary and run codegen again" && exit 1)
95100
working-directory: integ-tests/typescript
96101
- name: Ensure No Changes (2 of 3)
97-
run: git diff --exit-code
102+
run: |
103+
git diff --exit-code || (echo "::error:: merge canary and run codegen again" && exit 1)
98104
- name: Test Node Generator (3 of 3)
99-
run: pnpm generate
105+
run: |
106+
pnpm generate || (echo "::error:: merge canary and run codegen again" && exit 1)
100107
working-directory: integ-tests/typescript
101108
- name: Ensure No Changes (3 of 3)
102-
run: git diff --exit-code
109+
run: |
110+
git diff --exit-code || (echo "::error:: merge canary and run codegen again" && exit 1)
103111
build:
104112
runs-on: ubuntu-latest
105113
steps:

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -243,6 +243,7 @@ fn load_env_from_common_locations() -> Result<Option<(HashMap<String, String>, P
243243
Ok(None)
244244
}
245245

246+
// TODO: Is this even being used?
246247
/// Loads and applies environment variables to the current process
247248
pub fn dotenv(path: Option<PathBuf>) -> Result<Option<PathBuf>> {
248249
let (env_vars, path) = match path {

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

Lines changed: 67 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ use error::BamlError;
66
use indexmap::IndexMap;
77
use internal_baml_codegen::GeneratorArgs;
88
use json_response::Json;
9+
use std::collections::HashMap;
910

1011
use anyhow::{Context, Result};
1112
use arg_validation::BamlServeValidate;
@@ -58,6 +59,24 @@ pub struct ServeArgs {
5859
#[derive(Deserialize, Clone, Debug)]
5960
pub struct BamlOptions {
6061
pub client_registry: Option<ClientRegistry>,
62+
#[serde(default)]
63+
pub env: HashMap<String, Option<String>>,
64+
}
65+
66+
impl BamlOptions {
67+
fn env(&self) -> HashMap<String, String> {
68+
if self.env.is_empty() {
69+
return std::env::vars().collect();
70+
}
71+
let mut env = std::env::vars().collect::<HashMap<String, String>>();
72+
for (k, v) in &self.env {
73+
match v {
74+
Some(v) => env.insert(k.clone(), v.clone()),
75+
None => env.remove(k),
76+
};
77+
}
78+
env
79+
}
6180
}
6281

6382
impl ServeArgs {
@@ -311,12 +330,17 @@ Tip: test that the server is up using `curl http://localhost:{}/_debug/ping`
311330
Err(e) => return e.into_response(),
312331
};
313332

314-
let ctx_mgr = RuntimeContextManager::new_from_env_vars(std::env::vars().collect(), None);
315-
let client_registry = b_options.and_then(|options| options.client_registry);
333+
let client_registry = b_options.clone().and_then(|options| options.client_registry);
316334

317335
let locked = self.b.read().await;
336+
let env_vars: HashMap<String, String> = b_options
337+
.as_ref()
338+
.map_or_else(
339+
|| std::env::vars().collect(),
340+
|options| options.env(),
341+
);
318342
let (result, _trace_id) = locked
319-
.call_function(b_fn, &args, &ctx_mgr, None, client_registry.as_ref(), None)
343+
.call_function(b_fn, &args, &Default::default(), None, client_registry.as_ref(), None, env_vars)
320344
.await;
321345

322346
match result {
@@ -395,19 +419,24 @@ Tip: test that the server is up using `curl http://localhost:{}/_debug/ping`
395419
Err(e) => return e.into_response(),
396420
};
397421

398-
let client_registry = b_options.and_then(|options| options.client_registry);
422+
let client_registry = b_options.clone().and_then(|options| options.client_registry);
399423

400424
tokio::spawn(async move {
401-
let ctx_mgr =
402-
RuntimeContextManager::new_from_env_vars(std::env::vars().collect(), None);
425+
let env_vars: HashMap<String, String> = b_options
426+
.as_ref()
427+
.map_or_else(
428+
|| std::env::vars().collect(),
429+
|options| options.env(),
430+
);
403431

404432
let result_stream = self.b.read().await.stream_function(
405433
b_fn,
406434
&args,
407-
&ctx_mgr,
435+
&Default::default(),
408436
None,
409437
client_registry.as_ref(),
410438
Some(vec![]),
439+
env_vars,
411440
);
412441

413442
match result_stream {
@@ -424,9 +453,10 @@ Tip: test that the server is up using `curl http://localhost:{}/_debug/ping`
424453
}
425454
}
426455
}),
427-
&ctx_mgr,
456+
&Default::default(),
428457
None,
429458
None,
459+
HashMap::new(),
430460
)
431461
.await;
432462

@@ -730,4 +760,33 @@ mod tests {
730760
};
731761
assert_eq!(client_registry, expected_client_registry);
732762
}
763+
764+
#[test]
765+
fn test_parse_baml_options_with_env() {
766+
// Set up a dummy env var to test removal
767+
std::env::set_var("REMOVE_ME", "should_be_removed");
768+
std::env::set_var("KEEP_ME", "should_be_overwritten");
769+
770+
let baml_options: BamlOptions = serde_json::from_str(
771+
r#"
772+
{
773+
"env": {
774+
"NEW_VAR": "new_value",
775+
"KEEP_ME": "new_value",
776+
"REMOVE_ME": null
777+
}
778+
}"#,
779+
)
780+
.unwrap();
781+
782+
// The env() method should:
783+
// - Add NEW_VAR
784+
// - Overwrite KEEP_ME
785+
// - Remove REMOVE_ME
786+
let env_map = baml_options.env();
787+
788+
assert_eq!(env_map.get("NEW_VAR"), Some(&"new_value".to_string()));
789+
assert_eq!(env_map.get("KEEP_ME"), Some(&"new_value".to_string()));
790+
assert!(!env_map.contains_key("REMOVE_ME"));
791+
}
733792
}

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

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -247,9 +247,10 @@ async fn beta_reduce<'a>(
247247
tx.unbounded_send(vec![app_span]).unwrap();
248248
}
249249
if eval_final_llm_fn {
250+
// TODO: env vars are not supported yet for expressions.
250251
let res: anyhow::Result<FunctionResult> = env
251252
.runtime
252-
.call_function(name.clone(), &args_map, &ctx, None, None, None)
253+
.call_function(name.clone(), &args_map, &ctx, None, None, None, HashMap::new())
253254
.await
254255
.0;
255256

@@ -822,7 +823,7 @@ test Poems {
822823
let (res, _) = rt
823824
// .run_test("Second", "TestSecond", &ctx, Some(on_event))
824825
// .run_test("Go", "Go", &ctx, Some(on_event), None)
825-
.run_test("Poems", "Poems", &ctx, Some(on_event), None)
826+
.run_test("Poems", "Poems", &ctx, Some(on_event), None, HashMap::new())
826827
// .run_test("MakePerson", "TestMakePerson", &ctx, Some(on_event), None)
827828
// .run_test("CompareHaikus", "Test", &ctx, Some(on_event))
828829
// .run_test("LlmParseInt", "TestParse", &ctx, Some(on_event))
@@ -947,7 +948,7 @@ test TestMakePerson() {
947948
// dbg!(&f.item);
948949
let (res, _) = rt
949950
// .run_test("Second", "TestSecond", &ctx, Some(on_event))
950-
.run_test("OuterPyramid", "OuterPyramid", &ctx, Some(on_event), None)
951+
.run_test("OuterPyramid", "OuterPyramid", &ctx, Some(on_event), None, HashMap::new())
951952
// .run_test("MakePerson", "TestMakePerson", &ctx, Some(on_event), None)
952953
// .run_test("CompareHaikus", "Test", &ctx, Some(on_event))
953954
// .run_test("LlmParseInt", "TestParse", &ctx, Some(on_event))

‎engine/baml-runtime/src/internal/prompt_renderer/render_output_format.rs‎

Lines changed: 15 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -464,10 +464,12 @@ mod tests {
464464
)]
465465
.into_iter()
466466
.collect();
467-
let env_vars: HashMap<&str, &str> = HashMap::new();
468-
let baml_runtime = BamlRuntime::from_file_content(".", &files, env_vars).unwrap();
467+
let env_vars = HashMap::new();
468+
let baml_runtime = BamlRuntime::from_file_content(".", &files, env_vars.clone()).unwrap();
469469
let ctx_manager = baml_runtime.create_ctx_manager(BamlValue::Null, None);
470-
let ctx: RuntimeContext = ctx_manager.create_ctx(None, None, None).unwrap();
470+
let ctx: RuntimeContext = ctx_manager
471+
.create_ctx(None, None, env_vars.clone(), None)
472+
.unwrap();
471473

472474
let field_type = FieldType::Enum("Foo".to_string());
473475
let render_output =
@@ -521,10 +523,12 @@ class Resume {
521523
)]
522524
.into_iter()
523525
.collect();
524-
let env_vars: HashMap<&str, &str> = HashMap::new();
525-
let baml_runtime = BamlRuntime::from_file_content(".", &files, env_vars).unwrap();
526+
let env_vars = HashMap::new();
527+
let baml_runtime = BamlRuntime::from_file_content(".", &files, env_vars.clone()).unwrap();
526528
let ctx_manager = baml_runtime.create_ctx_manager(BamlValue::Null, None);
527-
let ctx: RuntimeContext = ctx_manager.create_ctx(None, None, None).unwrap();
529+
let ctx: RuntimeContext = ctx_manager
530+
.create_ctx(None, None, env_vars.clone(), None)
531+
.unwrap();
528532

529533
let field_type = FieldType::class("Resume");
530534
let render_output =
@@ -617,10 +621,12 @@ class Resume {
617621
)]
618622
.into_iter()
619623
.collect();
620-
let env_vars: HashMap<&str, &str> = HashMap::new();
621-
let baml_runtime = BamlRuntime::from_file_content(".", &files, env_vars).unwrap();
624+
let env_vars = HashMap::new();
625+
let baml_runtime = BamlRuntime::from_file_content(".", &files, env_vars.clone()).unwrap();
622626
let ctx_manager = baml_runtime.create_ctx_manager(BamlValue::Null, None);
623-
let ctx: RuntimeContext = ctx_manager.create_ctx(None, None, None).unwrap();
627+
let ctx: RuntimeContext = ctx_manager
628+
.create_ctx(None, None, env_vars.clone(), None)
629+
.unwrap();
624630

625631
let field_type = FieldType::class("Resume");
626632
let render_output =

0 commit comments

Comments
 (0)