esm compat (#1831)

- **esm compat**
- **Support esm modules**

<!-- ELLIPSIS_HIDDEN -->


----

> [!IMPORTANT]
> Adds ESM module support, modifies TypeScript import paths for ESM, and
includes tests and integration setup for ESM compatibility.
> 
>   - **ESM Support**:
> - Adds ESM module format support in `parse_generator()` in `v2.rs`.
> - Updates `generators.baml` to include `lang_typescript_esm` with
`module_format esm`.
>   - **TypeScript Imports**:
> - Adds `replace_ts_imports_with_js()` in `typescript/mod.rs` to modify
import paths to `.js` for ESM.
> - Uses `modify_files()` in `dir_writer.rs` to apply import path
changes.
>   - **Testing**:
> - Adds tests for `replace_ts_imports_with_js()` in
`typescript/mod.rs`.
> - Adds `typescript-esm` integration test setup with `package.json`,
`main.ts`, and `README.md`.
> 
> <sup>This description was created by </sup>[<img alt="Ellipsis"
src="https://img.shields.io/badge/Ellipsis-blue?color=175173">](https://www.ellipsis.dev?ref=BoundaryML%2Fbaml&utm_source=github&utm_medium=referral)<sup>
for c3dd2b4. You can
[customize](https://app.ellipsis.dev/BoundaryML/settings/summaries) this
summary. It will automatically update as commits are pushed.</sup>


<!-- ELLIPSIS_HIDDEN -->

Author: aaronvg

engine/Cargo.lock

@@ -34,6 +34,12 @@ pub enum Generator {
     BoundaryCloud(CloudProject),
 }
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum ModuleFormat {
+    Cjs,
+    Esm,
+}
+
 // TODO: we should figure out how to model generator fields using serde, since
 // the generator blocks are essentially a serde_json parse
 // problem is that serde_json has atrocious error messages and we need to provide
@@ -48,7 +54,8 @@ pub struct CodegenGenerator {
     output_dir: PathBuf,
     pub version: String,
     pub client_package_name: Option<String>,
-
+    // For TS generators, we can choose between CJS and ESM module formats
+    pub module_format: Option<ModuleFormat>,
     pub span: crate::ast::Span,
 }
 

engine/baml-lib/baml-core/src/configuration.rs

@@ -12,7 +12,7 @@ use strum::VariantNames;
 
 use crate::configuration::{
     CloudProject, CloudProjectBuilder, CodegenGeneratorBuilder, Generator,
-    GeneratorDefaultClientMode, GeneratorOutputType,
+    GeneratorDefaultClientMode, GeneratorOutputType, ModuleFormat,
 };
 
 fn parse_required_key<'a>(
@@ -200,6 +200,30 @@ pub(crate) fn parse_generator(
         }
     }
 
+    match parse_optional_key(&args, "module_format") {
+        Ok(Some("cjs")) => {
+            builder.module_format(Some(ModuleFormat::Cjs));
+        }
+        Ok(Some("esm")) => {
+            builder.module_format(Some(ModuleFormat::Esm));
+        }
+        Ok(Some(name)) => {
+            errors.push(DatamodelError::new_validation_error(
+                &format!("'{}' is not supported. Use one of: 'cjs' or 'esm'", name),
+                args.get("module_format")
+                    .map(|arg| arg.span().clone())
+                    .unwrap_or_else(|| ast_generator.span().clone()),
+            ));
+        }
+        Ok(None) => {
+            // TODO: add a warning if not set?
+            builder.module_format(None);
+        }
+        Err(err) => {
+            errors.push(err);
+        }
+    }
+
     if !errors.is_empty() {
         return Err(errors);
     }
@@ -298,6 +322,7 @@ fn check_property_allowlist<'ir>(
         "on_generate",
         "project",
         "client_package_name",
+        "module_format",
     ];
 
     let mut errors = vec![];

engine/baml-lib/baml-core/src/validate/generator_loader/v2.rs

@@ -91,6 +91,7 @@ impl GenerateArgs {
                         } else {
                             None
                         },
+                        None,
                     )
                     .context("Failed while resolving .baml paths in baml_src/")?,
                 )

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

@@ -574,6 +574,7 @@ Tip: test that the server is up using `curl http://localhost:{}/_debug/ping`
             Vec::new(),
             None,
             None,
+            None,
         )
         .map_err(|_| BamlError::InternalError {
             message: "Failed to make placeholder generator".to_string(),

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

@@ -855,6 +855,7 @@ impl BamlRuntime {
                         generator.on_generate.clone(),
                         Some(generator.output_type),
                         generator.client_package_name.clone(),
+                        generator.module_format,
                     )?,
                 ))
             })

engine/baml-runtime/src/lib.rs

@@ -24,6 +24,7 @@ internal-baml-core.workspace = true
 either.workspace = true
 env_logger.workspace = true
 log.workspace = true
+regex.workspace = true
 pathdiff = "0.1.0"
 serde.workspace = true
 serde_json.workspace = true

engine/language_client_codegen/Cargo.toml

@@ -126,6 +126,12 @@ impl<L: LanguageFeatures + Default> FileCollector<L> {
         );
     }
 
+    pub(super) fn modify_files(&mut self, mut modify: impl FnMut(&mut String)) {
+        for (_path, content) in self.files.iter_mut() {
+            modify(content);
+        }
+    }
+
     /// Ensure that a directory contains only files we generated before nuking it.
     ///
     /// This is a safety measure to prevent accidentally deleting user files.

engine/language_client_codegen/src/dir_writer.rs

@@ -282,6 +282,7 @@ class Foo {
             Vec::new(),
             Some(GeneratorOutputType::Go),
             Some("example.com/integ-tests".to_string()),
+            None,
         )
         .unwrap()
     }

engine/language_client_codegen/src/go/mod.rs

@@ -2,9 +2,10 @@ use anyhow::{Context, Result};
 use baml_types::{Constraint, ConstraintLevel, FieldType};
 use indexmap::IndexMap;
 use internal_baml_core::{
-    configuration::{GeneratorDefaultClientMode, GeneratorOutputType},
+    configuration::{GeneratorDefaultClientMode, GeneratorOutputType, ModuleFormat},
     ir::repr::IntermediateRepr,
 };
+use regex::Regex;
 use std::{
     collections::{BTreeMap, HashSet},
     path::{Path, PathBuf},
@@ -39,6 +40,9 @@ pub struct GeneratorArgs {
     // The type of client to generate
     client_type: Option<GeneratorOutputType>,
     client_package_name: Option<String>,
+
+    // For TS generators, we can choose between CJS and ESM module formats
+    module_format: Option<ModuleFormat>,
 }
 
 fn relative_path_to_baml_src(path: &Path, baml_src: &Path) -> Result<PathBuf> {
@@ -62,6 +66,7 @@ impl GeneratorArgs {
         on_generate: Vec<String>,
         client_type: Option<GeneratorOutputType>,
         client_package_name: Option<String>,
+        module_format: Option<ModuleFormat>,
     ) -> Result<Self> {
         let baml_src = baml_src_dir.into();
         let input_file_map: BTreeMap<PathBuf, String> = input_files
@@ -80,6 +85,7 @@ impl GeneratorArgs {
             on_generate,
             client_type,
             client_package_name,
+            module_format,
         })
     }
 

engine/language_client_codegen/src/lib.rs

@@ -464,6 +464,7 @@ class Foo {
             Vec::new(),
             Some(GeneratorOutputType::PythonPydantic),
             None,
+            None,
         )
         .unwrap()
     }

engine/language_client_codegen/src/python/mod.rs

@@ -11,9 +11,10 @@ use baml_types::LiteralValue;
 use generate_types::{render_docstring, type_name_for_checks};
 use indexmap::IndexMap;
 use internal_baml_core::{
-    configuration::{GeneratorDefaultClientMode, GeneratorOutputType},
+    configuration::{GeneratorDefaultClientMode, GeneratorOutputType, ModuleFormat},
     ir::{repr::IntermediateRepr, FieldType, IRHelper, IRHelperExtended},
 };
+use regex::Regex;
 
 use self::typescript_language_features::{ToTypescript, TypescriptLanguageFeatures};
 use crate::{dir_writer::FileCollector, field_type_attributes};
@@ -290,9 +291,60 @@ pub(crate) fn generate(
         TypescriptFramework::None => {}
     }
 
+    // if it's typescriopt and generator.esm is enabled, we need to change the imports in each file to use the .js extension
+    if generator.module_format == Some(ModuleFormat::Esm) {
+        baml_log::info!("Changing imports to .js for ESM");
+        collector.modify_files(|content: &mut String| {
+            *content = replace_ts_imports_with_js(content);
+        });
+    }
+
     collector.commit(&generator.output_dir())
 }
 
+fn replace_ts_imports_with_js(content: &str) -> String {
+    // Regex to find import/export statements with module specifiers.
+    // It captures the import/export part, quotes, and the path itself.
+    // Escaped curly braces in the character set just in case.
+    let re = Regex::new(r#"(import(?:["'\s]*(?:[\w\*\{\}\n\r\t, ]+)from\s*)?|export(?:["'\s]*(?:[\w\*\{\}\n\r\t, ]+)from\s*)?)(["'])([^"']+)(["'])"#).unwrap();
+
+    re.replace_all(content, |caps: &regex::Captures| {
+        let import_export_part = &caps[1];
+        let quote = &caps[2];
+        let path = &caps[3];
+        let closing_quote = &caps[4];
+
+        // Check if it's a relative path (starts with ./ or ../)
+        if path.starts_with("./") || path.starts_with("../") {
+            // Check if it already has a common JS/TS/CSS extension
+            if !path.ends_with(".js") &&
+               !path.ends_with(".mjs") &&
+               !path.ends_with(".cjs") &&
+               !path.ends_with(".jsx") && // Consider react specific extensions too
+               !path.ends_with(".tsx") &&
+               !path.ends_with(".css") && // Ignore CSS files
+               !path.ends_with(".json")
+            {
+                // Remove existing .ts if present before adding .js
+                let base_path = if path.ends_with(".ts") {
+                    &path[..path.len() - 3]
+                } else {
+                    path
+                };
+                // Append .js
+                format!("{import_export_part}{quote}{base_path}.js{closing_quote}")
+            } else {
+                // Already has a recognized extension, leave it as is.
+                caps[0].to_string()
+            }
+        } else {
+            // Not a relative path (e.g., external package like 'react' or '@boundaryml/baml'), leave it as is.
+            caps[0].to_string()
+        }
+    })
+    .to_string()
+}
+
 impl TryFrom<(&'_ IntermediateRepr, &'_ crate::GeneratorArgs)> for TypescriptConfig {
     type Error = anyhow::Error;
 
@@ -594,7 +646,9 @@ impl ToTypeReferenceInClientDefinition for FieldType {
             FieldType::WithMetadata { .. } => {
                 unreachable!("distribute_metadata makes this field unreachable.")
             }
-            FieldType::Arrow(_) => todo!("Arrow types should not be used in generated type definitions"),
+            FieldType::Arrow(_) => {
+                todo!("Arrow types should not be used in generated type definitions")
+            }
         };
         let base_type_ref = if is_partial_type {
             base_rep
@@ -675,7 +729,9 @@ impl ToTypeReferenceInClientDefinition for FieldType {
             FieldType::Optional(inner) => {
                 format!("{} | null", inner.to_type_ref(ir, use_module_prefix))
             }
-            FieldType::Arrow(_) => todo!("Arrow types should not be used in generated type definitions"),
+            FieldType::Arrow(_) => {
+                todo!("Arrow types should not be used in generated type definitions")
+            }
             FieldType::WithMetadata { base, .. } => match field_type_attributes(self) {
                 Some(checks) => {
                     let base_type_ref = base.to_type_ref(ir, use_module_prefix);
@@ -687,3 +743,99 @@ impl ToTypeReferenceInClientDefinition for FieldType {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_replace_ts_imports_with_js() {
+        // Add .js to relative paths without extension
+        assert_eq!(
+            replace_ts_imports_with_js("import { Foo } from './bar';"),
+            "import { Foo } from './bar.js';"
+        );
+        assert_eq!(
+            replace_ts_imports_with_js("export * from \"../baz/qux\";"),
+            "export * from \"../baz/qux.js\";"
+        );
+        assert_eq!(
+            replace_ts_imports_with_js("import type { Bar } from './bar'"),
+            "import type { Bar } from './bar.js'"
+        );
+        assert_eq!(
+            replace_ts_imports_with_js("import {\n  Thing1,\n  Thing2\n} from \"./things\";"),
+            "import {\n  Thing1,\n  Thing2\n} from \"./things.js\";"
+        );
+
+        // Replace .ts with .js in relative paths
+        assert_eq!(
+            replace_ts_imports_with_js("import { Foo } from './bar.ts';"),
+            "import { Foo } from './bar.js';"
+        );
+        assert_eq!(
+            replace_ts_imports_with_js("export * from \"../baz/qux.ts\";"),
+            "export * from \"../baz/qux.js\";"
+        );
+
+        // Should ignore already correct .js paths
+        assert_eq!(
+            replace_ts_imports_with_js("import { Foo } from './bar.js';"),
+            "import { Foo } from './bar.js';"
+        );
+        // Should ignore other extensions like .css, .mjs, .cjs
+        assert_eq!(
+            replace_ts_imports_with_js("import styles from './styles.css';"),
+            "import styles from './styles.css';"
+        );
+        assert_eq!(
+            replace_ts_imports_with_js("import config from './config.json';"),
+            "import config from './config.json';"
+        );
+        assert_eq!(
+            replace_ts_imports_with_js("import { util } from './util.mjs';"),
+            "import { util } from './util.mjs';"
+        );
+        assert_eq!(
+            replace_ts_imports_with_js("import { main } from '../main.cjs';"),
+            "import { main } from '../main.cjs';"
+        );
+        assert_eq!(
+            replace_ts_imports_with_js("import { Comp } from './Comp.tsx';"),
+            "import { Comp } from './Comp.tsx';"
+        );
+        assert_eq!(
+            replace_ts_imports_with_js("import { Button } from './Button.jsx';"),
+            "import { Button } from './Button.jsx';"
+        );
+
+        // Should ignore absolute paths or URLs
+        assert_eq!(
+            replace_ts_imports_with_js("import React from 'react';"),
+            "import React from 'react';"
+        );
+        assert_eq!(
+            replace_ts_imports_with_js("import { BamlClient } from '@boundaryml/baml';"),
+            "import { BamlClient } from '@boundaryml/baml';"
+        );
+        assert_eq!(
+            replace_ts_imports_with_js("const path = '/path/to/file.ts';"),
+            "const path = '/path/to/file.ts';" // This is not an import/export statement
+        );
+
+        // Empty string
+        assert_eq!(replace_ts_imports_with_js(""), "");
+        // String with no imports
+        assert_eq!(
+            replace_ts_imports_with_js("const x = 10; function y() {}"),
+            "const x = 10; function y() {}"
+        );
+        // Mixed content
+        assert_eq!(
+            replace_ts_imports_with_js(
+                "console.log('hello');\nimport { a } from './a.ts';\nimport { b } from './b';\nimport { c } from './c.js';\nimport { d } from 'd-lib';\nexport { e } from '../e.ts';\nconsole.log('world');"
+            ),
+            "console.log('hello');\nimport { a } from './a.js';\nimport { b } from './b.js';\nimport { c } from './c.js';\nimport { d } from 'd-lib';\nexport { e } from '../e.js';\nconsole.log('world');"
+        );
+    }
+}

engine/language_client_codegen/src/typescript/mod.rs

@@ -10,6 +10,15 @@ generator lang_typescript {
   version "0.84.4"
 }
 
+
+generator lang_typescript_esm {
+  output_type typescript
+  output_dir "../typescript-esm"
+  version "0.84.4"
+  module_format esm
+}
+
+
 generator lang_typescript_react {
   output_type typescript/react
   output_dir "../react"