diff --git a/engine/language-client-codegen/src/typescript/templates/tracing.ts.j2 b/engine/language-client-codegen/src/typescript/templates/tracing.ts.j2 index c79196ef6..275571706 100644 --- a/engine/language-client-codegen/src/typescript/templates/tracing.ts.j2 +++ b/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 } \ No newline at end of file diff --git a/engine/language-client-codegen/src/typescript/templates/type_builder.ts.j2 b/engine/language-client-codegen/src/typescript/templates/type_builder.ts.j2 index 703d0173c..4daa0b5ec 100644 --- a/engine/language-client-codegen/src/typescript/templates/type_builder.ts.j2 +++ b/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: Name): ClassBuilder { return this.tb.addClass(name); } diff --git a/engine/language_client_python/python_src/baml_py/baml_py.pyi b/engine/language_client_python/python_src/baml_py/baml_py.pyi index 31aec640b..797bb9b94 100644 --- a/engine/language_client_python/python_src/baml_py/baml_py.pyi +++ b/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: ... diff --git a/engine/language_client_python/python_src/baml_py/type_builder.py b/engine/language_client_python/python_src/baml_py/type_builder.py index 40f6abb30..f205bd903 100644 --- a/engine/language_client_python/python_src/baml_py/type_builder.py +++ b/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.") diff --git a/engine/language_client_python/src/types/type_builder.rs b/engine/language_client_python/src/types/type_builder.rs index a7ab16d7a..50485b221 100644 --- a/engine/language_client_python/src/types/type_builder.rs +++ b/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 { + let mut rs_types = vec![]; + for idx in 0..types.len() { + let item = types.get_item(idx)?; + let item = item.downcast::()?; + rs_types.push(item.borrow().inner.lock().unwrap().clone()); + } + Ok(baml_types::FieldType::union(rs_types).into()) + } } #[pymethods] diff --git a/engine/language_client_typescript/native.d.ts b/engine/language_client_typescript/native.d.ts index 723d589fc..259f02a14 100644 --- a/engine/language_client_typescript/native.d.ts +++ b/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 } export interface BamlLogEvent { diff --git a/engine/language_client_typescript/src/types/type_builder.rs b/engine/language_client_typescript/src/types/type_builder.rs index 007c182b2..dabfa5a8a 100644 --- a/engine/language_client_typescript/src/types/type_builder.rs +++ b/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] diff --git a/engine/language_client_typescript/type_builder.d.ts b/engine/language_client_typescript/type_builder.d.ts index 9a9c70b81..033c044ac 100644 --- a/engine/language_client_typescript/type_builder.d.ts +++ b/engine/language_client_typescript/type_builder.d.ts @@ -13,11 +13,14 @@ export declare class TypeBuilder { enums: Set; }); _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: Name, properties: Properties[]): ClassBuilder; enumBuilder(name: Name, values: T[]): EnumBuilder; addClass(name: Name): ClassBuilder; diff --git a/engine/language_client_typescript/type_builder.d.ts.map b/engine/language_client_typescript/type_builder.d.ts.map index c26a52d3a..fc281e5bd 100644 --- a/engine/language_client_typescript/type_builder.d.ts.map +++ b/engine/language_client_typescript/type_builder.d.ts.map @@ -1 +1 @@ -{"version":3,"file":"type_builder.d.ts","sourceRoot":"","sources":["typescript_src/type_builder.ts"],"names":[],"mappings":"AAAA,OAAO,EAGL,oBAAoB,IAAI,qBAAqB,EAC7C,gBAAgB,EAChB,SAAS,EACT,WAAW,IAAI,YAAY,EAC5B,MAAM,UAAU,CAAA;AAEjB,KAAK,SAAS,CAAC,CAAC,SAAS,MAAM,IAAI,MAAM,SAAS,CAAC,GAAG,KAAK,GAAG,IAAI,CAAA;AAClE,KAAK,MAAM,CAAC,CAAC,SAAS,MAAM,IAAI,SAAS,CAAC,CAAC,CAAC,SAAS,IAAI,GAAG,CAAC,GAAG,aAAa,CAAA;AAC7E,KAAK,UAAU,CAAC,CAAC,EAAE,QAAQ,SAAS,MAAM,EAAE,KAAK,SAAS,MAAM,IAAI,CAAC,CAAC,CAAC,SAAS,CAAC,KAAK,CAAC,GACnF,gCAAgC,KAAK,iCAAiC,MAAM,CAAC,QAAQ,CAAC,IAAI,GAC1F,CAAC,CAAA;AACL,KAAK,WAAW,CAAC,CAAC,EAAE,CAAC,IAAI,CAAC,SAAS,CAAC,GAAG,KAAK,GAAG,CAAC,CAAA;AAChD,KAAK,WAAW,CAAC,IAAI,SAAS,MAAM,EAAE,KAAK,SAAS,MAAM,EAAE,CAAC,SAAS,MAAM,IAAI,SAAS,CAAC,CAAC,CAAC,SAAS,IAAI,GACrG,UAAU,CAAC,WAAW,CAAC,KAAK,EAAE,CAAC,CAAC,EAAE,IAAI,EAAE,KAAK,CAAC,GAC9C,KAAK,CAAA;AAET,qBAAa,WAAW;IACtB,OAAO,CAAC,EAAE,CAAc;IACxB,SAAS,CAAC,OAAO,EAAE,GAAG,CAAC,MAAM,CAAC,CAAA;IAC9B,SAAS,CAAC,KAAK,EAAE,GAAG,CAAC,MAAM,CAAC,CAAA;gBAEhB,EAAE,OAAO,EAAE,KAAK,EAAE,EAAE;QAAE,OAAO,EAAE,GAAG,CAAC,MAAM,CAAC,CAAC;QAAC,KAAK,EAAE,GAAG,CAAC,MAAM,CAAC,CAAA;KAAE;IAM5E,GAAG,IAAI,YAAY;IAInB,MAAM,IAAI,SAAS;IAInB,GAAG,IAAI,SAAS;IAIhB,KAAK,IAAI,SAAS;IAIlB,IAAI,IAAI,SAAS;IAIjB,IAAI,CAAC,IAAI,EAAE,SAAS,GAAG,SAAS;IAIhC,YAAY,CAAC,IAAI,SAAS,MAAM,EAAE,UAAU,SAAS,MAAM,EACzD,IAAI,EAAE,IAAI,EACV,UAAU,EAAE,UAAU,EAAE,GACvB,YAAY,CAAC,IAAI,EAAE,UAAU,CAAC;IAIjC,WAAW,CAAC,IAAI,SAAS,MAAM,EAAE,CAAC,SAAS,MAAM,EAAE,IAAI,EAAE,IAAI,EAAE,MAAM,EAAE,CAAC,EAAE,GAAG,WAAW,CAAC,IAAI,EAAE,CAAC,CAAC;IAIjG,QAAQ,CAAC,IAAI,SAAS,MAAM,EAAE,IAAI,EAAE,IAAI,GAAG,YAAY,CAAC,IAAI,CAAC;IAW7D,OAAO,CAAC,IAAI,SAAS,MAAM,EAAE,IAAI,EAAE,IAAI,GAAG,WAAW,CAAC,IAAI,CAAC;CAU5D;AAED,qBAAa,YAAY,CAAC,SAAS,SAAS,MAAM,EAAE,UAAU,SAAS,MAAM,GAAG,MAAM;IAMlF,OAAO,CAAC,UAAU;IALpB,OAAO,CAAC,IAAI,CAAe;gBAGzB,EAAE,EAAE,YAAY,EAChB,IAAI,EAAE,SAAS,EACP,UAAU,GAAE,GAAG,CAAC,UAAU,GAAG,MAAM,CAAa;IAK1D,IAAI,IAAI,SAAS;IAIjB,cAAc,IAAI,KAAK,CAAC,CAAC,MAAM,EAAE,oBAAoB,CAAC,CAAC;IAIvD,WAAW,CAAC,CAAC,SAAS,MAAM,EAAE,IAAI,EAAE,WAAW,CAAC,SAAS,EAAE,CAAC,EAAE,UAAU,CAAC,EAAE,IAAI,EAAE,SAAS,GAAG,oBAAoB;IAQjH,QAAQ,CAAC,IAAI,EAAE,MAAM,GAAG,oBAAoB;CAM7C;AAED,cAAM,oBAAoB;IACxB,OAAO,CAAC,IAAI,CAAuB;gBAEvB,IAAI,EAAE,qBAAqB;IAIvC,KAAK,CAAC,KAAK,EAAE,MAAM,GAAG,IAAI,GAAG,oBAAoB;IAKjD,WAAW,CAAC,WAAW,EAAE,MAAM,GAAG,IAAI,GAAG,oBAAoB;CAI9D;AAED,qBAAa,WAAW,CAAC,QAAQ,SAAS,MAAM,EAAE,CAAC,SAAS,MAAM,GAAG,MAAM;IAMvE,OAAO,CAAC,MAAM;IALhB,OAAO,CAAC,IAAI,CAAc;gBAGxB,EAAE,EAAE,YAAY,EAChB,IAAI,EAAE,QAAQ,EACN,MAAM,GAAE,GAAG,CAAC,CAAC,GAAG,MAAM,CAAa;IAK7C,IAAI,IAAI,SAAS;IAIjB,KAAK,CAAC,CAAC,SAAS,MAAM,EAAE,IAAI,EAAE,CAAC,GAAG,CAAC,GAAG,gBAAgB;IAOtD,UAAU,IAAI,KAAK,CAAC,CAAC,MAAM,EAAE,gBAAgB,CAAC,CAAC;IAI/C,QAAQ,CAAC,CAAC,SAAS,MAAM,EAAE,IAAI,EAAE,WAAW,CAAC,QAAQ,EAAE,CAAC,EAAE,CAAC,CAAC,GAAG,gBAAgB;CAOhF"} \ No newline at end of file +{"version":3,"file":"type_builder.d.ts","sourceRoot":"","sources":["typescript_src/type_builder.ts"],"names":[],"mappings":"AAAA,OAAO,EAGL,oBAAoB,IAAI,qBAAqB,EAC7C,gBAAgB,EAChB,SAAS,EACT,WAAW,IAAI,YAAY,EAC5B,MAAM,UAAU,CAAA;AAEjB,KAAK,SAAS,CAAC,CAAC,SAAS,MAAM,IAAI,MAAM,SAAS,CAAC,GAAG,KAAK,GAAG,IAAI,CAAA;AAClE,KAAK,MAAM,CAAC,CAAC,SAAS,MAAM,IAAI,SAAS,CAAC,CAAC,CAAC,SAAS,IAAI,GAAG,CAAC,GAAG,aAAa,CAAA;AAC7E,KAAK,UAAU,CAAC,CAAC,EAAE,QAAQ,SAAS,MAAM,EAAE,KAAK,SAAS,MAAM,IAAI,CAAC,CAAC,CAAC,SAAS,CAAC,KAAK,CAAC,GACnF,gCAAgC,KAAK,iCAAiC,MAAM,CAAC,QAAQ,CAAC,IAAI,GAC1F,CAAC,CAAA;AACL,KAAK,WAAW,CAAC,CAAC,EAAE,CAAC,IAAI,CAAC,SAAS,CAAC,GAAG,KAAK,GAAG,CAAC,CAAA;AAChD,KAAK,WAAW,CAAC,IAAI,SAAS,MAAM,EAAE,KAAK,SAAS,MAAM,EAAE,CAAC,SAAS,MAAM,IAAI,SAAS,CAAC,CAAC,CAAC,SAAS,IAAI,GACrG,UAAU,CAAC,WAAW,CAAC,KAAK,EAAE,CAAC,CAAC,EAAE,IAAI,EAAE,KAAK,CAAC,GAC9C,KAAK,CAAA;AAET,qBAAa,WAAW;IACtB,OAAO,CAAC,EAAE,CAAc;IACxB,SAAS,CAAC,OAAO,EAAE,GAAG,CAAC,MAAM,CAAC,CAAA;IAC9B,SAAS,CAAC,KAAK,EAAE,GAAG,CAAC,MAAM,CAAC,CAAA;gBAEhB,EAAE,OAAO,EAAE,KAAK,EAAE,EAAE;QAAE,OAAO,EAAE,GAAG,CAAC,MAAM,CAAC,CAAC;QAAC,KAAK,EAAE,GAAG,CAAC,MAAM,CAAC,CAAA;KAAE;IAM5E,GAAG,IAAI,YAAY;IAInB,IAAI,IAAI,SAAS;IAIjB,MAAM,IAAI,SAAS;IAInB,GAAG,IAAI,SAAS;IAIhB,KAAK,IAAI,SAAS;IAIlB,IAAI,IAAI,SAAS;IAIjB,IAAI,CAAC,IAAI,EAAE,SAAS,GAAG,SAAS;IAIhC,GAAG,CAAC,OAAO,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,GAAG,SAAS;IAIxD,KAAK,CAAC,KAAK,EAAE,SAAS,EAAE,GAAG,SAAS;IAIpC,YAAY,CAAC,IAAI,SAAS,MAAM,EAAE,UAAU,SAAS,MAAM,EACzD,IAAI,EAAE,IAAI,EACV,UAAU,EAAE,UAAU,EAAE,GACvB,YAAY,CAAC,IAAI,EAAE,UAAU,CAAC;IAIjC,WAAW,CAAC,IAAI,SAAS,MAAM,EAAE,CAAC,SAAS,MAAM,EAAE,IAAI,EAAE,IAAI,EAAE,MAAM,EAAE,CAAC,EAAE,GAAG,WAAW,CAAC,IAAI,EAAE,CAAC,CAAC;IAIjG,QAAQ,CAAC,IAAI,SAAS,MAAM,EAAE,IAAI,EAAE,IAAI,GAAG,YAAY,CAAC,IAAI,CAAC;IAW7D,OAAO,CAAC,IAAI,SAAS,MAAM,EAAE,IAAI,EAAE,IAAI,GAAG,WAAW,CAAC,IAAI,CAAC;CAU5D;AAED,qBAAa,YAAY,CAAC,SAAS,SAAS,MAAM,EAAE,UAAU,SAAS,MAAM,GAAG,MAAM;IAMlF,OAAO,CAAC,UAAU;IALpB,OAAO,CAAC,IAAI,CAAe;gBAGzB,EAAE,EAAE,YAAY,EAChB,IAAI,EAAE,SAAS,EACP,UAAU,GAAE,GAAG,CAAC,UAAU,GAAG,MAAM,CAAa;IAK1D,IAAI,IAAI,SAAS;IAIjB,cAAc,IAAI,KAAK,CAAC,CAAC,MAAM,EAAE,oBAAoB,CAAC,CAAC;IAIvD,WAAW,CAAC,CAAC,SAAS,MAAM,EAAE,IAAI,EAAE,WAAW,CAAC,SAAS,EAAE,CAAC,EAAE,UAAU,CAAC,EAAE,IAAI,EAAE,SAAS,GAAG,oBAAoB;IAQjH,QAAQ,CAAC,IAAI,EAAE,MAAM,GAAG,oBAAoB;CAM7C;AAED,cAAM,oBAAoB;IACxB,OAAO,CAAC,IAAI,CAAuB;gBAEvB,IAAI,EAAE,qBAAqB;IAIvC,KAAK,CAAC,KAAK,EAAE,MAAM,GAAG,IAAI,GAAG,oBAAoB;IAKjD,WAAW,CAAC,WAAW,EAAE,MAAM,GAAG,IAAI,GAAG,oBAAoB;CAI9D;AAED,qBAAa,WAAW,CAAC,QAAQ,SAAS,MAAM,EAAE,CAAC,SAAS,MAAM,GAAG,MAAM;IAMvE,OAAO,CAAC,MAAM;IALhB,OAAO,CAAC,IAAI,CAAc;gBAGxB,EAAE,EAAE,YAAY,EAChB,IAAI,EAAE,QAAQ,EACN,MAAM,GAAE,GAAG,CAAC,CAAC,GAAG,MAAM,CAAa;IAK7C,IAAI,IAAI,SAAS;IAIjB,KAAK,CAAC,CAAC,SAAS,MAAM,EAAE,IAAI,EAAE,CAAC,GAAG,CAAC,GAAG,gBAAgB;IAOtD,UAAU,IAAI,KAAK,CAAC,CAAC,MAAM,EAAE,gBAAgB,CAAC,CAAC;IAI/C,QAAQ,CAAC,CAAC,SAAS,MAAM,EAAE,IAAI,EAAE,WAAW,CAAC,QAAQ,EAAE,CAAC,EAAE,CAAC,CAAC,GAAG,gBAAgB;CAOhF"} \ No newline at end of file diff --git a/engine/language_client_typescript/type_builder.js b/engine/language_client_typescript/type_builder.js index bde0e4e18..e4f75f63e 100644 --- a/engine/language_client_typescript/type_builder.js +++ b/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)); } diff --git a/engine/language_client_typescript/typescript_src/type_builder.ts b/engine/language_client_typescript/typescript_src/type_builder.ts index 471eada0e..d300b0ab2 100644 --- a/engine/language_client_typescript/typescript_src/type_builder.ts +++ b/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: Name, properties: Properties[], diff --git a/integ-tests/python/poetry.lock b/integ-tests/python/poetry.lock index 7a6417105..0a5d5450c 100644 --- a/integ-tests/python/poetry.lock +++ b/integ-tests/python/poetry.lock @@ -329,6 +329,17 @@ files = [ {file = "tomli-2.0.1.tar.gz", hash = "sha256:de526c12914f0c550d15924c62d72abc48d6fe7364aa87328337a31007fe8a4f"}, ] +[[package]] +name = "types-assertpy" +version = "1.1.0.20240712" +description = "Typing stubs for assertpy" +optional = false +python-versions = ">=3.8" +files = [ + {file = "types-assertpy-1.1.0.20240712.tar.gz", hash = "sha256:c332d09298b158f09890af12f14d83ebb9f0e5014e91219ae145a6a5af9220f9"}, + {file = "types_assertpy-1.1.0.20240712-py3-none-any.whl", hash = "sha256:3962d4d863e4c091b9a8d7f02e3ee0d2ac06f25bdca0e6e4f848cf54c7214455"}, +] + [[package]] name = "typing-extensions" version = "4.12.2" @@ -343,4 +354,4 @@ files = [ [metadata] lock-version = "2.0" python-versions = "^3.8" -content-hash = "75bc9294db7074ba805acc12ba1a8a201c8099a40f8ffe1d2e2aefb14d48e5e6" +content-hash = "0f05e76164bce73b7880e1de6c9a7105771fda78f3bd70241e59a85e95027f40" diff --git a/integ-tests/python/pyproject.toml b/integ-tests/python/pyproject.toml index 8c58d850b..6588c977d 100644 --- a/integ-tests/python/pyproject.toml +++ b/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" diff --git a/integ-tests/python/tests/test_functions.py b/integ-tests/python/tests/test_functions.py index b00cb9f66..ba7b1660d 100644 --- a/integ-tests/python/tests/test_functions.py +++ b/integ-tests/python/tests/test_functions.py @@ -1,4 +1,5 @@ import time +from typing import List import pytest from assertpy import assert_that from dotenv import load_dotenv @@ -11,11 +12,12 @@ from ..baml_client.globals import ( DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_RUNTIME, ) +from ..baml_client import partial_types from ..baml_client.types import ( DynInputOutput, NamedArgsSingleEnumList, NamedArgsSingleClass, - StringToClassEntry, + StringToClassEntry, ) from ..baml_client.tracing import trace, set_tags, flush, on_log_event from ..baml_client.type_builder import TypeBuilder @@ -25,7 +27,6 @@ import random - def test_sync(): res = sync_b.TestFnNamedArgsSingleClass( myArg=NamedArgsSingleClass( @@ -94,18 +95,22 @@ async def test_single_int(self): @pytest.mark.asyncio async def test_single_map_string_to_string(self): - res = await b.TestFnNamedArgsSingleMapStringToString({'lorem': 'ipsum', 'dolor': 'sit'}) - assert 'lorem' in res + res = await b.TestFnNamedArgsSingleMapStringToString( + {"lorem": "ipsum", "dolor": "sit"} + ) + assert "lorem" in res @pytest.mark.asyncio async def test_single_map_string_to_class(self): - res = await b.TestFnNamedArgsSingleMapStringToClass({'lorem': StringToClassEntry(word='ipsum')}) - assert res['lorem'].word == 'ipsum' + res = await b.TestFnNamedArgsSingleMapStringToClass( + {"lorem": StringToClassEntry(word="ipsum")} + ) + assert res["lorem"].word == "ipsum" @pytest.mark.asyncio async def test_single_map_string_to_map(self): - res = await b.TestFnNamedArgsSingleMapStringToMap({'lorem': {'word': 'ipsum'}}) - assert res['lorem']['word'] == 'ipsum' + res = await b.TestFnNamedArgsSingleMapStringToMap({"lorem": {"word": "ipsum"}}) + assert res["lorem"]["word"] == "ipsum" class MyCustomClass(NamedArgsSingleClass): @@ -418,7 +423,7 @@ async def dummy_fn(foo: str): def test_tracing_sync(): # res = parent_sync("first-arg-value") - res2 = sync_dummy_func("second-dummycall-arg") + _ = sync_dummy_func("second-dummycall-arg") def test_tracing_thread_pool(): @@ -491,7 +496,6 @@ def parent_sync(myStr: str): return "hello world parentsync" - @trace async def async_dummy_func(myArgggg: str): await asyncio.sleep(0.5 + random.random()) @@ -504,7 +508,7 @@ def sync_dummy_func(dummyFuncArg: str): @pytest.fixture(scope="session", autouse=True) -def cleanup(request): +def cleanup(): """Cleanup a testing directory once we are finished.""" flush() @@ -600,7 +604,7 @@ async def test_dynamic_class_nested_output_stream(): input="My name is Mark Gonzalez. My hair is black and I'm 6 feet tall.", baml_options={"tb": tb}, ) - msgs = [] + msgs: List[partial_types.DynamicOutput] = [] async for msg in stream: print("streamed ", msg) print("streamed ", msg.model_dump()) @@ -630,7 +634,7 @@ async def test_stream_dynamic_class_output(): input="My name is Harrison. My hair is black and I'm 6 feet tall.", baml_options={"tb": tb, "client_registry": cr}, ) - msgs = [] + msgs: List[partial_types.DynamicOutput] = [] async for msg in stream: print("streamed ", msg.model_dump()) msgs.append(msg) @@ -653,8 +657,8 @@ async def test_dynamic_inputs_list2(): res = await b.DynamicListInputOutput( [ - DynInputOutput( - **{ + DynInputOutput.model_validate( + { "new_key": "hi1", "testKey": "myTest", "blah": { @@ -662,13 +666,15 @@ async def test_dynamic_inputs_list2(): }, } ), - { - "new_key": "hi", - "testKey": "myTest", - "blah": { - "nestedKey1": "nestedVal", - }, - }, + DynInputOutput.model_validate( + { + "new_key": "hi", + "testKey": "myTest", + "blah": { + "nestedKey1": "nestedVal", + }, + } + ), ], {"tb": tb}, ) @@ -690,20 +696,24 @@ async def test_dynamic_inputs_list(): res = await b.DynamicListInputOutput( [ - { - "new_key": "hi", - "testKey": "myTest", - "blah": { - "nestedKey1": "nestedVal", - }, - }, - { - "new_key": "hi", - "testKey": "myTest", - "blah": { - "nestedKey1": "nestedVal", - }, - }, + DynInputOutput.model_validate( + { + "new_key": "hi", + "testKey": "myTest", + "blah": { + "nestedKey1": "nestedVal", + }, + } + ), + DynInputOutput.model_validate( + { + "new_key": "hi", + "testKey": "myTest", + "blah": { + "nestedKey1": "nestedVal", + }, + } + ), ], {"tb": tb}, ) @@ -715,12 +725,84 @@ async def test_dynamic_inputs_list(): assert res[1].blah["nestedKey1"] == "nestedVal" +@pytest.mark.asyncio +async def test_dynamic_output_map(): + tb = TypeBuilder() + tb.DynamicOutput.add_property("hair_color", tb.string()) + tb.DynamicOutput.add_property( + "attributes", tb.map(tb.string(), tb.string()) + ).description("Things like 'eye_color' or 'facial_hair'") + print(tb.DynamicOutput.list_properties()) + for prop, _ in tb.DynamicOutput.list_properties(): + print(f"Property: {prop}") + + res = await b.MyFunc( + input="My name is Harrison. My hair is black and I'm 6 feet tall. I have blue eyes and a beard.", + baml_options={"tb": tb}, + ) + + print("final ", res) + print("final ", res.model_dump()) + print("final ", res.model_dump_json()) + assert res.hair_color == "black" + assert res.attributes["eye_color"] == "blue" + assert res.attributes["facial_hair"] == "beard" + + +@pytest.mark.asyncio +async def test_dynamic_output_union(): + tb = TypeBuilder() + tb.DynamicOutput.add_property("hair_color", tb.string()) + tb.DynamicOutput.add_property( + "attributes", tb.map(tb.string(), tb.string()) + ).description("Things like 'eye_color' or 'facial_hair'") + # Define two classes + class1 = tb.add_class("Class1") + class1.add_property("meters", tb.float()) + + class2 = tb.add_class("Class2") + class2.add_property("feet", tb.float()) + class2.add_property("inches", tb.float().optional()) + + # Use the classes in a union property + tb.DynamicOutput.add_property("height", tb.union([class1.type(), class2.type()])) + print(tb.DynamicOutput.list_properties()) + for prop, _ in tb.DynamicOutput.list_properties(): + print(f"Property: {prop}") + + res = await b.MyFunc( + input="My name is Harrison. My hair is black and I'm 6 feet tall. I have blue eyes and a beard. I am 30 years old.", + baml_options={"tb": tb}, + ) + + print("final ", res) + print("final ", res.model_dump()) + print("final ", res.model_dump_json()) + assert res.hair_color == "black" + assert res.attributes["eye_color"] == "blue" + assert res.attributes["facial_hair"] == "beard" + assert res.height["feet"] == 6 + + res = await b.MyFunc( + input="My name is Harrison. My hair is black and I'm 1.8 meters tall. I have blue eyes and a beard. I am 30 years old.", + baml_options={"tb": tb}, + ) + + print("final ", res) + print("final ", res.model_dump()) + print("final ", res.model_dump_json()) + assert res.hair_color == "black" + assert res.attributes["eye_color"] == "blue" + assert res.attributes["facial_hair"] == "beard" + assert res.height["meters"] == 1.8 + + @pytest.mark.asyncio async def test_nested_class_streaming(): stream = b.stream.FnOutputClassNested( input="My name is Harrison. My hair is black and I'm 6 feet tall." ) - msgs = [] + msgs: List[partial_types.TestClassNested] = [] async for msg in stream: print("streamed ", msg.model_dump(mode="json")) msgs.append(msg) @@ -791,7 +873,7 @@ async def test_stream_serialization_exception(): async for msg in stream: print("streamed ", msg) - res = await stream.get_final_response() + _ = await stream.get_final_response() print("Exception message: ", excinfo) assert "Failed to coerce" in str(excinfo) @@ -807,7 +889,7 @@ async def stream_func(): async for msg in stream: print("streamed ", msg) - res = await stream.get_final_response() + _ = await stream.get_final_response() print("Exception message: ", excinfo) assert "Failed to coerce" in str(excinfo) diff --git a/integ-tests/typescript/baml_client/tracing.ts b/integ-tests/typescript/baml_client/tracing.ts index db3cc2deb..75638f124 100644 --- a/integ-tests/typescript/baml_client/tracing.ts +++ b/integ-tests/typescript/baml_client/tracing.ts @@ -27,7 +27,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 } \ No newline at end of file diff --git a/integ-tests/typescript/baml_client/type_builder.ts b/integ-tests/typescript/baml_client/type_builder.ts index aa1463576..680d13bc9 100644 --- a/integ-tests/typescript/baml_client/type_builder.ts +++ b/integ-tests/typescript/baml_client/type_builder.ts @@ -126,6 +126,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: Name): ClassBuilder { return this.tb.addClass(name); } diff --git a/integ-tests/typescript/tests/integ-tests.test.ts b/integ-tests/typescript/tests/integ-tests.test.ts index bcc76d26c..dced62fac 100644 --- a/integ-tests/typescript/tests/integ-tests.test.ts +++ b/integ-tests/typescript/tests/integ-tests.test.ts @@ -421,6 +421,75 @@ describe('Integ tests', () => { expect(res[0]['testKey']).toEqual('myTest') }) + it('should work with dynamic output map', async () => { + let tb = new TypeBuilder() + tb.DynamicOutput.addProperty('hair_color', tb.string()) + tb.DynamicOutput.addProperty('attributes', tb.map(tb.string(), tb.string())).description( + "Things like 'eye_color' or 'facial_hair'", + ) + console.log(tb.DynamicOutput.listProperties()) + for (const [prop, _] of tb.DynamicOutput.listProperties()) { + console.log(`Property: ${prop}`) + } + + const res = await b.MyFunc( + "My name is Harrison. My hair is black and I'm 6 feet tall. I have blue eyes and a beard.", + { tb }, + ) + + console.log('final ', res) + console.log('final ', res.modelDump()) + console.log('final ', res.modelDumpJson()) + expect(res.hair_color).toEqual('black') + expect(res.attributes['eye_color']).toEqual('blue') + expect(res.attributes['facial_hair']).toEqual('beard') + }) + + it('should work with dynamic output union', async () => { + let tb = new TypeBuilder() + tb.DynamicOutput.addProperty('hair_color', tb.string()) + tb.DynamicOutput.addProperty('attributes', tb.map(tb.string(), tb.string())).description( + "Things like 'eye_color' or 'facial_hair'", + ) + + // Define two classes + const class1 = tb.addClass('Class1') + class1.addProperty('meters', tb.float()) + + const class2 = tb.addClass('Class2') + class2.addProperty('feet', tb.float()) + class2.addProperty('inches', tb.float().optional()) + + // Use the classes in a union property + tb.DynamicOutput.addProperty('height', tb.union([class1.type(), class2.type()])) + console.log(tb.DynamicOutput.listProperties()) + for (const [prop, _] of tb.DynamicOutput.listProperties()) { + console.log(`Property: ${prop}`) + } + + let res = await b.MyFunc( + "My name is Harrison. My hair is black and I'm 6 feet tall. I have blue eyes and a beard. I am 30 years old.", + { tb }, + ) + + console.log('final ', res) + expect(res.hair_color).toEqual('black') + expect(res.attributes['eye_color']).toEqual('blue') + expect(res.attributes['facial_hair']).toEqual('beard') + expect(res.height['feet']).toEqual(6) + + res = await b.MyFunc( + "My name is Harrison. My hair is black and I'm 1.8 meters tall. I have blue eyes and a beard. I am 30 years old.", + { tb }, + ) + + console.log('final ', res) + expect(res.hair_color).toEqual('black') + expect(res.attributes['eye_color']).toEqual('blue') + expect(res.attributes['facial_hair']).toEqual('beard') + expect(res.height['meters']).toEqual(1.8) + }) + // test with extra list, boolean in the input as well. it('should work with nested classes', async () => {