Add support for unions / maps / null in TypeBuilder. (#820)

* Minor bug fix: allow unregistration of onLogEvent hook in typescript

engine/language-client-codegen/src/typescript/templates/tracing.ts.j2

@@ -10,7 +10,7 @@ DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_CTX.upsertTags.bind(DO_NOT_
 const flush = () => {
   DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_CTX.flush.bind(DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_CTX)()
 }
-const onLogEvent = (callback: (event: BamlLogEvent) => void) =>
+const onLogEvent = (callback: undefined | (event: BamlLogEvent) => void) =>
 DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_CTX.onLogEvent(callback)
 
 export { traceAsync, traceSync, setTags, flush, onLogEvent }

engine/language-client-codegen/src/typescript/templates/type_builder.ts.j2

@@ -57,6 +57,18 @@ export default class TypeBuilder {
         return this.tb.list(type)
     }
 
+    null(): FieldType {
+        return this.tb.null()
+    }
+
+    map(key: FieldType, value: FieldType): FieldType {
+        return this.tb.map(key, value)
+    }
+
+    union(types: FieldType[]): FieldType {
+        return this.tb.union(types)
+    }
+
     addClass<Name extends string>(name: Name): ClassBuilder<Name> {
         return this.tb.addClass(name);
     }

engine/language_client_python/python_src/baml_py/baml_py.pyi

@@ -1,4 +1,4 @@
-from typing import Any, Callable, Dict, Optional, Tuple
+from typing import Any, Callable, Dict, List, Optional, Tuple
 
 class FunctionResult:
     """The result of a BAML function call.
@@ -171,6 +171,8 @@ class TypeBuilder:
     def list(self, element_type: FieldType) -> FieldType: ...
     def null(self) -> FieldType: ...
     def optional(self, inner_type: FieldType) -> FieldType: ...
+    def map(self, key_type: FieldType, value_type: FieldType) -> FieldType: ...
+    def union(self, *types: FieldType) -> FieldType: ...
 
 class ClientRegistry:
     def __init__(self) -> None: ...

engine/language_client_python/python_src/baml_py/type_builder.py

@@ -34,6 +34,15 @@ def bool(self):
     def list(self, inner: FieldType):
         return self._tb.list(inner)
 
+    def null(self):
+        return self._tb.null()
+
+    def map(self, key: FieldType, value: FieldType):
+        return self._tb.map(key, value)
+
+    def union(self, types: typing.List[FieldType]):
+        return self._tb.union(*types)
+
     def add_class(self, name: str) -> "NewClassBuilder":
         if name in self.__classes:
             raise ValueError(f"Class with name {name} already exists.")

engine/language_client_python/src/types/type_builder.rs

@@ -1,6 +1,11 @@
 use baml_runtime::type_builder::{self, WithMeta};
 use baml_types::BamlValue;
-use pyo3::pymethods;
+use pyo3::{
+    prelude::PyAnyMethods,
+    pymethods,
+    types::{PyTuple, PyTupleMethods},
+    Bound, PyResult,
+};
 
 crate::lang_wrapper!(TypeBuilder, type_builder::TypeBuilder);
 crate::lang_wrapper!(EnumBuilder, type_builder::EnumBuilder, sync_thread_safe, name: String);
@@ -67,6 +72,25 @@ impl TypeBuilder {
     pub fn null(&self) -> FieldType {
         baml_types::FieldType::null().into()
     }
+
+    pub fn map(&self, key: &FieldType, value: &FieldType) -> FieldType {
+        baml_types::FieldType::map(
+            key.inner.lock().unwrap().clone(),
+            value.inner.lock().unwrap().clone(),
+        )
+        .into()
+    }
+
+    #[pyo3(signature = (*types))]
+    pub fn union<'py>(&self, types: &Bound<'_, PyTuple>) -> PyResult<FieldType> {
+        let mut rs_types = vec![];
+        for idx in 0..types.len() {
+            let item = types.get_item(idx)?;
+            let item = item.downcast::<FieldType>()?;
+            rs_types.push(item.borrow().inner.lock().unwrap().clone());
+        }
+        Ok(baml_types::FieldType::union(rs_types).into())
+    }
 }
 
 #[pymethods]

engine/language_client_typescript/native.d.ts

@@ -106,6 +106,8 @@ export class TypeBuilder {
   float(): FieldType
   bool(): FieldType
   null(): FieldType
+  map(key: FieldType, value: FieldType): FieldType
+  union(types: Array<FieldType>): FieldType
 }
 
 export interface BamlLogEvent {

engine/language_client_typescript/src/types/type_builder.rs

@@ -74,6 +74,26 @@ impl TypeBuilder {
     pub fn null(&self) -> FieldType {
         baml_types::FieldType::null().into()
     }
+
+    #[napi]
+    pub fn map(&self, key: &FieldType, value: &FieldType) -> FieldType {
+        baml_types::FieldType::map(
+            key.inner.lock().unwrap().clone(),
+            value.inner.lock().unwrap().clone(),
+        )
+        .into()
+    }
+
+    #[napi]
+    pub fn union(&self, types: Vec<&FieldType>) -> FieldType {
+        baml_types::FieldType::union(
+            types
+                .iter()
+                .map(|t| t.inner.lock().unwrap().clone())
+                .collect(),
+        )
+        .into()
+    }
 }
 
 #[napi]

engine/language_client_typescript/type_builder.d.ts

@@ -13,11 +13,14 @@ export declare class TypeBuilder {
         enums: Set<string>;
     });
     _tb(): _TypeBuilder;
+    null(): FieldType;
     string(): FieldType;
     int(): FieldType;
     float(): FieldType;
     bool(): FieldType;
     list(type: FieldType): FieldType;
+    map(keyType: FieldType, valueType: FieldType): FieldType;
+    union(types: FieldType[]): FieldType;
     classBuilder<Name extends string, Properties extends string>(name: Name, properties: Properties[]): ClassBuilder<Name, Properties>;
     enumBuilder<Name extends string, T extends string>(name: Name, values: T[]): EnumBuilder<Name, T>;
     addClass<Name extends string>(name: Name): ClassBuilder<Name>;

engine/language_client_typescript/type_builder.js

@@ -14,6 +14,9 @@ class TypeBuilder {
     _tb() {
         return this.tb;
     }
+    null() {
+        return this.tb.null();
+    }
     string() {
         return this.tb.string();
     }
@@ -29,6 +32,12 @@ class TypeBuilder {
     list(type) {
         return this.tb.list(type);
     }
+    map(keyType, valueType) {
+        return this.tb.map(keyType, valueType);
+    }
+    union(types) {
+        return this.tb.union(types);
+    }
     classBuilder(name, properties) {
         return new ClassBuilder(this.tb, name, new Set(properties));
     }

engine/language_client_typescript/typescript_src/type_builder.ts

@@ -32,6 +32,10 @@ export class TypeBuilder {
     return this.tb
   }
 
+  null(): FieldType {
+    return this.tb.null()
+  }
+
   string(): FieldType {
     return this.tb.string()
   }
@@ -52,6 +56,14 @@ export class TypeBuilder {
     return this.tb.list(type)
   }
 
+  map(keyType: FieldType, valueType: FieldType): FieldType {
+    return this.tb.map(keyType, valueType)
+  }
+
+  union(types: FieldType[]): FieldType {
+    return this.tb.union(types)
+  }
+
   classBuilder<Name extends string, Properties extends string>(
     name: Name,
     properties: Properties[],

integ-tests/python/pyproject.toml

@@ -21,6 +21,9 @@ pydantic = "^2.7.1"
 python-dotenv = "^1.0.1"
 assertpy = "^1.1"
 
+[tool.poetry.group.dev.dependencies]
+types-assertpy = "^1.1.0.20240712"
+
 [build-system]
 requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"