diff --git a/.gitignore b/.gitignore index 15e8215d2..a82d4d444 100644 --- a/.gitignore +++ b/.gitignore @@ -122,6 +122,10 @@ $RECYCLE.BIN/ /dist /node_modules /out/ +engine/language_client_typescript/*.d.ts +engine/language_client_typescript/*.d.ts.map +engine/language_client_typescript/*.js +!engine/language_client_typescript/cli.js engine/language_client_ruby/**/*.bundle engine/target/ Cargo.lock diff --git a/engine/.gitignore b/engine/.gitignore deleted file mode 100644 index a7456499a..000000000 --- a/engine/.gitignore +++ /dev/null @@ -1,5 +0,0 @@ -target -*.env -**/__pycache__ -test_logs/** -*.so \ No newline at end of file diff --git a/engine/baml-runtime/src/lib.rs b/engine/baml-runtime/src/lib.rs index a6cbc5660..b01844bbf 100644 --- a/engine/baml-runtime/src/lib.rs +++ b/engine/baml-runtime/src/lib.rs @@ -327,7 +327,7 @@ impl ExperimentalTracingInterface for BamlRuntime { } } - fn flush(&self) -> Result<()> { + fn flush(&self) -> Result { self.tracer.flush() } diff --git a/engine/baml-runtime/src/runtime_interface.rs b/engine/baml-runtime/src/runtime_interface.rs index 2d05e4e33..4d3b7c1a8 100644 --- a/engine/baml-runtime/src/runtime_interface.rs +++ b/engine/baml-runtime/src/runtime_interface.rs @@ -92,7 +92,7 @@ pub trait ExperimentalTracingInterface { ctx: &RuntimeContextManager, ) -> Result>; - fn flush(&self) -> Result<()>; + fn flush(&self) -> Result; #[cfg(not(target_arch = "wasm32"))] fn set_log_event_callback(&self, callback: LogEventCallbackSync) -> Result<()>; diff --git a/engine/baml-runtime/src/tracing/mod.rs b/engine/baml-runtime/src/tracing/mod.rs index c4e925c25..f22bb5502 100644 --- a/engine/baml-runtime/src/tracing/mod.rs +++ b/engine/baml-runtime/src/tracing/mod.rs @@ -1,22 +1,20 @@ pub mod api_wrapper; -#[cfg(not(target_arch = "wasm32"))] -mod threaded_tracer; -#[cfg(target_arch = "wasm32")] -mod wasm_tracer; use crate::on_log_event::LogEventCallbackSync; -use anyhow::Result; +use anyhow::{Context, Result}; use baml_types::{BamlMap, BamlMediaType, BamlValue}; +use cfg_if::cfg_if; use colored::Colorize; use internal_baml_jinja::RenderedPrompt; use std::collections::HashMap; +use std::sync::{Arc, Mutex}; use uuid::Uuid; use crate::{ internal::llm_client::LLMResponse, tracing::api_wrapper::core_types::Role, type_builder::TypeBuilder, FunctionResult, RuntimeContext, RuntimeContextManager, SpanCtx, - TestResponse, + TestResponse, TraceStats, }; use self::api_wrapper::{ @@ -27,16 +25,17 @@ use self::api_wrapper::{ }, APIWrapper, }; -#[cfg(not(target_arch = "wasm32"))] -use self::threaded_tracer::ThreadedTracer; -#[cfg(target_arch = "wasm32")] -use self::wasm_tracer::NonThreadedTracer; +cfg_if! { + if #[cfg(target_arch = "wasm32")] { + mod wasm_tracer; + use self::wasm_tracer::NonThreadedTracer as TracerImpl; + } else { + mod threaded_tracer; + use self::threaded_tracer::ThreadedTracer as TracerImpl; + } +} -#[cfg(not(target_arch = "wasm32"))] -type TracerImpl = ThreadedTracer; -#[cfg(target_arch = "wasm32")] -type TracerImpl = NonThreadedTracer; #[derive(Debug)] pub struct TracingSpan { span_id: Uuid, @@ -48,6 +47,15 @@ pub struct BamlTracer { options: APIWrapper, enabled: bool, tracer: Option, + trace_stats: Arc>, +} + +macro_rules! bail { + ($self:ident, $($arg:tt)*) => { + let mut locked = $self.trace_stats.lock().unwrap(); + locked.n_spans_failed_before_submit += 1; + anyhow::bail!($($arg)*); + }; } impl BamlTracer { @@ -68,6 +76,9 @@ impl BamlTracer { }, enabled: options.enabled(), options, + trace_stats: Arc::new(Mutex::new(TraceStats { + n_spans_failed_before_submit: 0, + })), }; tracer } @@ -79,12 +90,12 @@ impl BamlTracer { } } - pub(crate) fn flush(&self) -> Result<()> { - if let Some(tracer) = &self.tracer { - tracer.flush() - } else { - Ok(()) + pub(crate) fn flush(&self) -> Result { + if let Some(ref tracer) = self.tracer { + tracer.flush().context("Failed to flush BAML traces")?; } + + Ok(self.trace_stats.lock().unwrap().clone()) } pub(crate) fn start_span( @@ -95,6 +106,7 @@ impl BamlTracer { params: &BamlMap, ) -> (Option, RuntimeContext) { let span_id = ctx.enter(function_name); + log::trace!("Entering span {:#?} in {:?}", span_id, function_name); if !self.enabled { return (None, ctx.create_ctx(tb)); } @@ -115,7 +127,8 @@ impl BamlTracer { response: Option, ) -> Result> { let Some((span_id, event_chain, tags)) = ctx.exit() else { - anyhow::bail!( + bail!( + self, "Attempting to finish a span {:#?} without first starting one. Current context {:#?}", span, ctx @@ -123,7 +136,7 @@ impl BamlTracer { }; if span.span_id != span_id { - anyhow::bail!("Span ID mismatch: {} != {}", span.span_id, span_id); + bail!(self, "Span ID mismatch: {} != {}", span.span_id, span_id); } if let Some(tracer) = &self.tracer { @@ -144,20 +157,37 @@ impl BamlTracer { response: Option, ) -> Result> { let Some((span_id, event_chain, tags)) = ctx.exit() else { - anyhow::bail!( + bail!(self, "Attempting to finish a span {:#?} without first starting one. Current context {:#?}", span, ctx ); }; + log::trace!( + "Finishing span: {:#?} {}\nevent chain {:?}", + span, + span_id, + event_chain + ); if span.span_id != span_id { - anyhow::bail!("Span ID mismatch: {} != {}", span.span_id, span_id); + bail!(self, "Span ID mismatch: {} != {}", span.span_id, span_id); } if let Some(tracer) = &self.tracer { - tracer.submit(response.to_log_schema(&self.options, event_chain, tags, span))?; - Ok(Some(span_id)) + cfg_if! { + if #[cfg(target_arch = "wasm32")] { + async { + tracer + .submit(response.to_log_schema(&self.options, event_chain, tags, span)) + .await?; + Ok(Some(span_id)) + } + } else { + tracer.submit(response.to_log_schema(&self.options, event_chain, tags, span))?; + Ok(Some(span_id)) + } + } } else { Ok(None) } @@ -171,11 +201,14 @@ impl BamlTracer { response: &Result, ) -> Result> { let Some((span_id, event_chain, tags)) = ctx.exit() else { - anyhow::bail!("Attempting to finish a span without first starting one"); + bail!( + self, + "Attempting to finish a span without first starting one" + ); }; if span.span_id != span_id { - anyhow::bail!("Span ID mismatch: {} != {}", span.span_id, span_id); + bail!(self, "Span ID mismatch: {} != {}", span.span_id, span_id); } if let Ok(response) = &response { @@ -208,11 +241,21 @@ impl BamlTracer { response: &Result, ) -> Result> { let Some((span_id, event_chain, tags)) = ctx.exit() else { - anyhow::bail!("Attempting to finish a span without first starting one"); + bail!( + self, + "Attempting to finish a span without first starting one" + ); }; + log::trace!( + "Finishing baml span: {:#?} {}\nevent chain {:?}", + span, + span_id, + event_chain + ); + if span.span_id != span_id { - anyhow::bail!("Span ID mismatch: {} != {}", span.span_id, span_id); + bail!(self, "Span ID mismatch: {} != {}", span.span_id, span_id); } if let Ok(response) = &response { diff --git a/engine/baml-runtime/src/tracing/threaded_tracer.rs b/engine/baml-runtime/src/tracing/threaded_tracer.rs index ff3b0433b..41c3a6e33 100644 --- a/engine/baml-runtime/src/tracing/threaded_tracer.rs +++ b/engine/baml-runtime/src/tracing/threaded_tracer.rs @@ -1,10 +1,7 @@ -use std::{ - sync::mpsc::{Receiver, Sender, TryRecvError}, -}; - -// use crate::log_callback_event::LogEvent use anyhow::Result; -use web_time::Duration; +use std::sync::mpsc::{channel, Receiver, RecvTimeoutError, Sender, TryRecvError}; +use std::sync::{Arc, Mutex}; +use web_time::{Duration, Instant}; use crate::{ on_log_event::{LogEvent, LogEventCallbackSync, LogEventMetadata}, @@ -54,7 +51,7 @@ fn batch_processor( ) { let api_config = &api_config; let mut batch = Vec::with_capacity(max_batch_size); - let mut now = std::time::Instant::now(); + let mut now = Instant::now(); let rt = tokio::runtime::Runtime::new().unwrap(); loop { // Try to fill the batch up to max_batch_size @@ -65,8 +62,8 @@ fn batch_processor( } Ok(TxEventSignal::Flush) => (false, true, false), Ok(TxEventSignal::Stop) => (false, false, true), - Err(std::sync::mpsc::RecvTimeoutError::Timeout) => (false, false, false), - Err(std::sync::mpsc::RecvTimeoutError::Disconnected) => (false, false, true), + Err(RecvTimeoutError::Timeout) => (false, false, false), + Err(RecvTimeoutError::Disconnected) => (false, false, true), }; let time_trigger = now.elapsed().as_millis() >= 1000; @@ -100,11 +97,11 @@ fn batch_processor( pub(super) struct ThreadedTracer { api_config: APIWrapper, - tx: std::sync::Arc>>, - rx: std::sync::Arc>>, + tx: Arc>>, + rx: Arc>>, #[allow(dead_code)] join_handle: std::thread::JoinHandle<()>, - log_event_callback: std::sync::Arc>>, + log_event_callback: Arc>>, } impl ThreadedTracer { @@ -112,12 +109,12 @@ impl ThreadedTracer { api_config: APIWrapper, max_batch_size: usize, ) -> ( - std::sync::mpsc::Sender, - std::sync::mpsc::Receiver, + Sender, + Receiver, std::thread::JoinHandle<()>, ) { - let (tx, rx) = std::sync::mpsc::channel(); - let (stop_tx, stop_rx) = std::sync::mpsc::channel(); + let (tx, rx) = channel(); + let (stop_tx, stop_rx) = channel(); let join_handle = std::thread::spawn(move || batch_processor(api_config, rx, stop_tx, max_batch_size)); @@ -128,10 +125,10 @@ impl ThreadedTracer { let (tx, rx, join_handle) = Self::start_worker(api_config.clone(), max_batch_size); Self { api_config: api_config.clone(), - tx: std::sync::Arc::new(std::sync::Mutex::new(tx)), - rx: std::sync::Arc::new(std::sync::Mutex::new(rx)), + tx: Arc::new(Mutex::new(tx)), + rx: Arc::new(Mutex::new(rx)), join_handle, - log_event_callback: std::sync::Arc::new(std::sync::Mutex::new(None)), + log_event_callback: Arc::new(Mutex::new(None)), } } diff --git a/engine/baml-runtime/src/types/mod.rs b/engine/baml-runtime/src/types/mod.rs index c7eefb381..e29dddd17 100644 --- a/engine/baml-runtime/src/types/mod.rs +++ b/engine/baml-runtime/src/types/mod.rs @@ -4,8 +4,10 @@ pub mod on_log_event; mod response; pub(crate) mod runtime_context; mod stream; +mod trace_stats; pub use context_manager::RuntimeContextManager; pub use response::{FunctionResult, TestFailReason, TestResponse, TestStatus}; pub use runtime_context::{RuntimeContext, SpanCtx}; pub use stream::FunctionResultStream; +pub use trace_stats::TraceStats; diff --git a/engine/baml-runtime/src/types/trace_stats.rs b/engine/baml-runtime/src/types/trace_stats.rs new file mode 100644 index 000000000..99af656aa --- /dev/null +++ b/engine/baml-runtime/src/types/trace_stats.rs @@ -0,0 +1,15 @@ +/// Stats about all the spans sent to the tracer. +/// +/// A span has the following lifecycle and can fail at any of these points: +/// +/// ```text +/// start -> finalize (ctx.exit) -> submit -> send +/// ``` +/// +/// TODO: return stats about the # of spans successfully sent +#[derive(Clone)] +pub struct TraceStats { + /// # of spans that we called finish_span or finish_baml_span on + /// but did not submit due to an error + pub n_spans_failed_before_submit: u32, +} diff --git a/engine/language-client-codegen/src/python/templates/tracing.py.j2 b/engine/language-client-codegen/src/python/templates/tracing.py.j2 index 07a25cf84..6c132c8e0 100644 --- a/engine/language-client-codegen/src/python/templates/tracing.py.j2 +++ b/engine/language-client-codegen/src/python/templates/tracing.py.j2 @@ -2,7 +2,8 @@ from .globals import DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_CTX trace = DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_CTX.trace_fn set_tags = DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_CTX.upsert_tags -flush = DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_CTX.flush +def flush(): + DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_CTX.flush() on_log_event = DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_CTX.on_log_event diff --git a/engine/language-client-codegen/src/typescript/generate_types.rs b/engine/language-client-codegen/src/typescript/generate_types.rs index 8038be2b9..8f67d1693 100644 --- a/engine/language-client-codegen/src/typescript/generate_types.rs +++ b/engine/language-client-codegen/src/typescript/generate_types.rs @@ -7,14 +7,14 @@ use crate::GeneratorArgs; use super::ToTypeReferenceInClientDefinition; #[derive(askama::Template)] -#[template(path = "type_builder.js.j2", escape = "none")] +#[template(path = "type_builder.ts.j2", escape = "none")] pub(crate) struct TypeBuilder<'ir> { enums: Vec>, classes: Vec>, } #[derive(askama::Template)] -#[template(path = "types.js.j2", escape = "none")] +#[template(path = "types.ts.j2", escape = "none")] pub(crate) struct TypescriptTypes<'ir> { enums: Vec>, classes: Vec>, diff --git a/engine/language-client-codegen/src/typescript/mod.rs b/engine/language-client-codegen/src/typescript/mod.rs index 64da56fa3..00b5226f4 100644 --- a/engine/language-client-codegen/src/typescript/mod.rs +++ b/engine/language-client-codegen/src/typescript/mod.rs @@ -12,7 +12,7 @@ use self::typescript_language_features::{ToTypescript, TypescriptLanguageFeature use crate::dir_writer::FileCollector; #[derive(askama::Template)] -#[template(path = "client.js.j2", escape = "none")] +#[template(path = "client.ts.j2", escape = "none")] struct TypescriptClient { funcs: Vec, types: Vec, diff --git a/engine/language-client-codegen/src/typescript/templates/client.js.j2 b/engine/language-client-codegen/src/typescript/templates/client.ts.j2 similarity index 87% rename from engine/language-client-codegen/src/typescript/templates/client.js.j2 rename to engine/language-client-codegen/src/typescript/templates/client.ts.j2 index 4982c5809..4157808ca 100644 --- a/engine/language-client-codegen/src/typescript/templates/client.js.j2 +++ b/engine/language-client-codegen/src/typescript/templates/client.ts.j2 @@ -11,9 +11,13 @@ export type RecursivePartialNull = T extends object : T | null; export class BamlClient { + private runtime: BamlRuntime + private ctx_manager: BamlCtxManager private stream_client: BamlStreamClient - constructor(private runtime: BamlRuntime, private ctx_manager: BamlCtxManager) { + constructor(runtime: BamlRuntime, ctx_manager: BamlCtxManager) { + this.runtime = runtime + this.ctx_manager = ctx_manager this.stream_client = new BamlStreamClient(runtime, ctx_manager) } @@ -35,7 +39,7 @@ export class BamlClient { "{{name}}": {{name}}{% if optional %}?? null{% endif %}{% if !loop.last %},{% endif %} {%- endfor %} }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as {{fn.return_type}} @@ -61,14 +65,14 @@ class BamlStreamClient { {%- endfor %} }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, {{ fn.return_type }}>( raw, (a): a is RecursivePartialNull<{{ fn.return_type }}> => a, (a): a is {{ fn.return_type }} => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } 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 7523436df..c79196ef6 100644 --- a/engine/language-client-codegen/src/typescript/templates/tracing.ts.j2 +++ b/engine/language-client-codegen/src/typescript/templates/tracing.ts.j2 @@ -2,13 +2,14 @@ import { BamlLogEvent } from '@boundaryml/baml'; import { DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_CTX } from './globals'; const traceAsync = -DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_CTX.traceFnAync.bind(DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_CTX) +DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_CTX.traceFnAsync.bind(DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_CTX) const traceSync = DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_CTX.traceFnSync.bind(DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_CTX) const setTags = DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_CTX.upsertTags.bind(DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_CTX) -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 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) => DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_CTX.onLogEvent(callback) diff --git a/engine/language-client-codegen/src/typescript/templates/type_builder.js.j2 b/engine/language-client-codegen/src/typescript/templates/type_builder.ts.j2 similarity index 100% rename from engine/language-client-codegen/src/typescript/templates/type_builder.js.j2 rename to engine/language-client-codegen/src/typescript/templates/type_builder.ts.j2 diff --git a/engine/language-client-codegen/src/typescript/templates/types.js.j2 b/engine/language-client-codegen/src/typescript/templates/types.ts.j2 similarity index 100% rename from engine/language-client-codegen/src/typescript/templates/types.js.j2 rename to engine/language-client-codegen/src/typescript/templates/types.ts.j2 diff --git a/engine/language_client_python/src/lib.rs b/engine/language_client_python/src/lib.rs index 7ecf4bacb..4f280f559 100644 --- a/engine/language_client_python/src/lib.rs +++ b/engine/language_client_python/src/lib.rs @@ -1,4 +1,5 @@ mod parse_py_type; +mod runtime; mod types; use pyo3::prelude::{pyfunction, pymodule, PyAnyMethods, PyModule, PyResult}; @@ -33,7 +34,8 @@ fn baml_py(_: Python<'_>, m: Bound<'_, PyModule>) -> PyResult<()> { eprintln!("Failed to initialize BAML logger: {:#}", e); }; - m.add_class::()?; + m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; @@ -46,8 +48,9 @@ fn baml_py(_: Python<'_>, m: Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; - m.add_class::()?; - m.add_class::()?; + + m.add_class::()?; + m.add_class::()?; m.add_wrapped(wrap_pyfunction!(invoke_runtime_cli))?; diff --git a/engine/language_client_python/src/types/runtime.rs b/engine/language_client_python/src/runtime.rs similarity index 94% rename from engine/language_client_python/src/types/runtime.rs rename to engine/language_client_python/src/runtime.rs index 02d6d42de..af2fc35fd 100644 --- a/engine/language_client_python/src/types/runtime.rs +++ b/engine/language_client_python/src/runtime.rs @@ -1,10 +1,11 @@ use crate::parse_py_type::parse_py_type; use crate::types::function_results::FunctionResult; +use crate::types::trace_stats::TraceStats; use crate::BamlError; -use super::function_result_stream::FunctionResultStream; -use super::runtime_ctx_manager::RuntimeContextManager; -use super::type_builder::TypeBuilder; +use crate::types::function_result_stream::FunctionResultStream; +use crate::types::runtime_ctx_manager::RuntimeContextManager; +use crate::types::type_builder::TypeBuilder; use baml_runtime::runtime_interface::ExperimentalTracingInterface; use baml_runtime::BamlRuntime as CoreBamlRuntime; use pyo3::prelude::{pymethods, PyResult}; @@ -174,8 +175,11 @@ impl BamlRuntime { } #[pyo3()] - fn flush(&self) -> PyResult<()> { - self.inner.flush().map_err(BamlError::from_anyhow) + fn flush(&self) -> PyResult { + self.inner + .flush() + .map(Into::into) + .map_err(BamlError::from_anyhow) } #[pyo3()] diff --git a/engine/language_client_python/src/types/function_result_stream.rs b/engine/language_client_python/src/types/function_result_stream.rs index face146b2..cb868bc78 100644 --- a/engine/language_client_python/src/types/function_result_stream.rs +++ b/engine/language_client_python/src/types/function_result_stream.rs @@ -14,7 +14,7 @@ crate::lang_wrapper!( ); impl FunctionResultStream { - pub(super) fn new( + pub(crate) fn new( inner: baml_runtime::FunctionResultStream, event: Option, tb: Option, diff --git a/engine/language_client_python/src/types/mod.rs b/engine/language_client_python/src/types/mod.rs index 8c20c8c30..ab84069c3 100644 --- a/engine/language_client_python/src/types/mod.rs +++ b/engine/language_client_python/src/types/mod.rs @@ -1,20 +1,19 @@ mod lang_wrapper; -mod audio; -mod function_result_stream; -mod function_results; -mod image; -mod runtime; -mod runtime_ctx_manager; -mod span; -mod type_builder; +pub(crate) mod audio; +pub(crate) mod function_result_stream; +pub(crate) mod function_results; +pub(crate) mod image; +pub(crate) mod runtime_ctx_manager; +pub(crate) mod span; +pub(crate) mod trace_stats; +pub(crate) mod type_builder; pub use audio::BamlAudioPy; pub use function_result_stream::FunctionResultStream; pub use function_results::FunctionResult; pub use image::BamlImagePy; -pub use runtime::{BamlLogEvent, BamlRuntime, LogEventMetadata}; pub use runtime_ctx_manager::RuntimeContextManager; pub use span::BamlSpan; pub use type_builder::*; diff --git a/engine/language_client_python/src/types/span.rs b/engine/language_client_python/src/types/span.rs index 191631d84..f4a6e9494 100644 --- a/engine/language_client_python/src/types/span.rs +++ b/engine/language_client_python/src/types/span.rs @@ -6,8 +6,8 @@ use pyo3::{PyObject, Python, ToPyObject}; use crate::parse_py_type::parse_py_type; use crate::BamlError; -use super::runtime::BamlRuntime; use super::runtime_ctx_manager::RuntimeContextManager; +use crate::runtime::BamlRuntime; crate::lang_wrapper!(BamlSpan, Option>, diff --git a/engine/language_client_python/src/types/trace_stats.rs b/engine/language_client_python/src/types/trace_stats.rs new file mode 100644 index 000000000..502ee8479 --- /dev/null +++ b/engine/language_client_python/src/types/trace_stats.rs @@ -0,0 +1,10 @@ +use pyo3::pymethods; + +crate::lang_wrapper!(TraceStats, baml_runtime::TraceStats); + +#[pymethods] +impl TraceStats { + pub fn n_spans_failed_before_submit(&self) -> u32 { + self.inner.n_spans_failed_before_submit + } +} diff --git a/engine/language_client_typescript/async_context_vars.d.ts b/engine/language_client_typescript/async_context_vars.d.ts deleted file mode 100644 index 61cf8de62..000000000 --- a/engine/language_client_typescript/async_context_vars.d.ts +++ /dev/null @@ -1,16 +0,0 @@ -import { BamlSpan, RuntimeContextManager, BamlRuntime, BamlLogEvent } from './native'; -export declare class CtxManager { - private rt; - private ctx; - constructor(rt: BamlRuntime); - upsertTags(tags: Record): void; - get(): RuntimeContextManager; - startTraceSync(name: string, args: Record): BamlSpan; - startTraceAsync(name: string, args: Record): BamlSpan; - endTrace(span: BamlSpan, response: any): void; - flush(): void; - onLogEvent(callback: (event: BamlLogEvent) => void): void; - traceFnSync ReturnType>(name: string, func: F): F; - traceFnAync Promise>(name: string, func: F): F; -} -//# sourceMappingURL=async_context_vars.d.ts.map \ No newline at end of file diff --git a/engine/language_client_typescript/async_context_vars.d.ts.map b/engine/language_client_typescript/async_context_vars.d.ts.map deleted file mode 100644 index 3d7ae88ad..000000000 --- a/engine/language_client_typescript/async_context_vars.d.ts.map +++ /dev/null @@ -1 +0,0 @@ -{"version":3,"file":"async_context_vars.d.ts","sourceRoot":"","sources":["typescript_src/async_context_vars.ts"],"names":[],"mappings":"AAAA,OAAO,EAAE,QAAQ,EAAE,qBAAqB,EAAE,WAAW,EAAE,YAAY,EAAE,MAAM,UAAU,CAAA;AAGrF,qBAAa,UAAU;IACrB,OAAO,CAAC,EAAE,CAAa;IACvB,OAAO,CAAC,GAAG,CAA0C;gBAEzC,EAAE,EAAE,WAAW;IAS3B,UAAU,CAAC,IAAI,EAAE,MAAM,CAAC,MAAM,EAAE,MAAM,CAAC,GAAG,IAAI;IAK9C,GAAG,IAAI,qBAAqB;IAS5B,cAAc,CAAC,IAAI,EAAE,MAAM,EAAE,IAAI,EAAE,MAAM,CAAC,MAAM,EAAE,GAAG,CAAC,GAAG,QAAQ;IAOjE,eAAe,CAAC,IAAI,EAAE,MAAM,EAAE,IAAI,EAAE,MAAM,CAAC,MAAM,EAAE,GAAG,CAAC,GAAG,QAAQ;IAOlE,QAAQ,CAAC,IAAI,EAAE,QAAQ,EAAE,QAAQ,EAAE,GAAG,GAAG,IAAI;IAS7C,KAAK,IAAI,IAAI;IAIb,UAAU,CAAC,QAAQ,EAAE,CAAC,KAAK,EAAE,YAAY,KAAK,IAAI,GAAG,IAAI;IAQzD,WAAW,CAAC,UAAU,EAAE,CAAC,SAAS,CAAC,GAAG,IAAI,EAAE,GAAG,EAAE,KAAK,UAAU,EAAE,IAAI,EAAE,MAAM,EAAE,IAAI,EAAE,CAAC,GAAG,CAAC;IAsB3F,WAAW,CAAC,UAAU,EAAE,CAAC,SAAS,CAAC,GAAG,IAAI,EAAE,GAAG,EAAE,KAAK,OAAO,CAAC,UAAU,CAAC,EAAE,IAAI,EAAE,MAAM,EAAE,IAAI,EAAE,CAAC,GAAG,CAAC;CAqBrG"} \ No newline at end of file diff --git a/engine/language_client_typescript/async_context_vars.js b/engine/language_client_typescript/async_context_vars.js deleted file mode 100644 index 2c1a77f0d..000000000 --- a/engine/language_client_typescript/async_context_vars.js +++ /dev/null @@ -1,97 +0,0 @@ -"use strict"; -Object.defineProperty(exports, "__esModule", { value: true }); -exports.CtxManager = void 0; -const native_1 = require("./native"); -const async_hooks_1 = require("async_hooks"); -class CtxManager { - rt; - ctx; - constructor(rt) { - this.rt = rt; - this.ctx = new async_hooks_1.AsyncLocalStorage(); - this.ctx.enterWith(rt.createContextManager()); - process.on('exit', () => { - this.rt.flush(); - }); - } - upsertTags(tags) { - const manager = this.ctx.getStore(); - manager.upsertTags(tags); - } - get() { - let store = this.ctx.getStore(); - if (store === undefined) { - store = this.rt.createContextManager(); - this.ctx.enterWith(store); - } - return store; - } - startTraceSync(name, args) { - const mng = this.get(); - // const clone = mng.deepClone() - // this.ctx.enterWith(clone) - return native_1.BamlSpan.new(this.rt, name, args, mng); - } - startTraceAsync(name, args) { - const mng = this.get(); - const clone = mng.deepClone(); - this.ctx.enterWith(clone); - return native_1.BamlSpan.new(this.rt, name, args, clone); - } - endTrace(span, response) { - const manager = this.ctx.getStore(); - if (!manager) { - console.error('Context lost before span could be finished\n'); - return; - } - span.finish(response, manager); - } - flush() { - this.rt.flush(); - } - onLogEvent(callback) { - this.rt.setLogEventCallback((error, param) => { - if (!error) { - callback(param); - } - }); - } - traceFnSync(name, func) { - return ((...args) => { - const params = args.reduce((acc, arg, i) => ({ - ...acc, - [`arg${i}`]: arg, // generic way to label args - }), {}); - const span = this.startTraceSync(name, params); - try { - const response = func(...args); - this.endTrace(span, response); - return response; - } - catch (e) { - this.endTrace(span, e); - throw e; - } - }); - } - traceFnAync(name, func) { - const funcName = name; - return (async (...args) => { - const params = args.reduce((acc, arg, i) => ({ - ...acc, - [`arg${i}`]: arg, // generic way to label args - }), {}); - const span = this.startTraceAsync(funcName, params); - try { - const response = await func(...args); - this.endTrace(span, response); - return response; - } - catch (e) { - this.endTrace(span, e); - throw e; - } - }); - } -} -exports.CtxManager = CtxManager; diff --git a/engine/language_client_typescript/index.d.ts b/engine/language_client_typescript/index.d.ts deleted file mode 100644 index cb2886eea..000000000 --- a/engine/language_client_typescript/index.d.ts +++ /dev/null @@ -1,4 +0,0 @@ -export { BamlRuntime, FunctionResult, FunctionResultStream, BamlImage as Image, BamlAudio as Audio, invoke_runtime_cli } from './native'; -export { BamlStream } from './stream'; -export { CtxManager as BamlCtxManager } from './async_context_vars'; -//# sourceMappingURL=index.d.ts.map \ No newline at end of file diff --git a/engine/language_client_typescript/index.d.ts.map b/engine/language_client_typescript/index.d.ts.map deleted file mode 100644 index bc7bb3b91..000000000 --- a/engine/language_client_typescript/index.d.ts.map +++ /dev/null @@ -1 +0,0 @@ -{"version":3,"file":"index.d.ts","sourceRoot":"","sources":["typescript_src/index.ts"],"names":[],"mappings":"AAAA,OAAO,EAAE,WAAW,EAAE,cAAc,EAAE,oBAAoB,EAAE,SAAS,IAAI,KAAK,EAAE,SAAS,IAAI,KAAK,EAAE,kBAAkB,EAAE,MAAM,UAAU,CAAA;AACxI,OAAO,EAAE,UAAU,EAAE,MAAM,UAAU,CAAA;AACrC,OAAO,EAAE,UAAU,IAAI,cAAc,EAAE,MAAM,sBAAsB,CAAA"} \ No newline at end of file diff --git a/engine/language_client_typescript/index.js b/engine/language_client_typescript/index.js deleted file mode 100644 index 4379795f9..000000000 --- a/engine/language_client_typescript/index.js +++ /dev/null @@ -1,14 +0,0 @@ -"use strict"; -Object.defineProperty(exports, "__esModule", { value: true }); -exports.BamlCtxManager = exports.BamlStream = exports.invoke_runtime_cli = exports.Audio = exports.Image = exports.FunctionResultStream = exports.FunctionResult = exports.BamlRuntime = void 0; -var native_1 = require("./native"); -Object.defineProperty(exports, "BamlRuntime", { enumerable: true, get: function () { return native_1.BamlRuntime; } }); -Object.defineProperty(exports, "FunctionResult", { enumerable: true, get: function () { return native_1.FunctionResult; } }); -Object.defineProperty(exports, "FunctionResultStream", { enumerable: true, get: function () { return native_1.FunctionResultStream; } }); -Object.defineProperty(exports, "Image", { enumerable: true, get: function () { return native_1.BamlImage; } }); -Object.defineProperty(exports, "Audio", { enumerable: true, get: function () { return native_1.BamlAudio; } }); -Object.defineProperty(exports, "invoke_runtime_cli", { enumerable: true, get: function () { return native_1.invoke_runtime_cli; } }); -var stream_1 = require("./stream"); -Object.defineProperty(exports, "BamlStream", { enumerable: true, get: function () { return stream_1.BamlStream; } }); -var async_context_vars_1 = require("./async_context_vars"); -Object.defineProperty(exports, "BamlCtxManager", { enumerable: true, get: function () { return async_context_vars_1.CtxManager; } }); diff --git a/engine/language_client_typescript/native.d.ts b/engine/language_client_typescript/native.d.ts deleted file mode 100644 index 494513884..000000000 --- a/engine/language_client_typescript/native.d.ts +++ /dev/null @@ -1,106 +0,0 @@ -/* auto-generated by NAPI-RS */ -/* eslint-disable */ -export class BamlAudio { - static fromUrl(url: string): BamlAudio - static fromBase64(mediaType: string, base64: string): BamlAudio - isUrl(): boolean - asUrl(): string - asBase64(): [string, string] - toJSON(): any -} - -export class BamlImage { - static fromUrl(url: string): BamlImage - static fromBase64(mediaType: string, base64: string): BamlImage - isUrl(): boolean - asUrl(): string - asBase64(): [string, string] - toJSON(): any -} - -export class BamlRuntime { - static fromDirectory(directory: string, envVars: Record): BamlRuntime - static fromFiles(rootPath: string, files: Record, envVars: Record): BamlRuntime - createContextManager(): RuntimeContextManager - callFunction(functionName: string, args: { [string]: any }, ctx: RuntimeContextManager, tb?: TypeBuilder | undefined | null): Promise - streamFunction(functionName: string, args: { [string]: any }, cb: (err: any, param: FunctionResult) => void, ctx: RuntimeContextManager, tb?: TypeBuilder | undefined | null): FunctionResultStream - setLogEventCallback(func: (err: any, param: BamlLogEvent) => void): void - flush(): void -} - -export class BamlSpan { - static new(runtime: BamlRuntime, functionName: string, args: any, ctx: RuntimeContextManager): BamlSpan - finish(result: any, ctx: RuntimeContextManager): any -} - -export class ClassBuilder { - field(): FieldType - property(name: string): ClassPropertyBuilder -} - -export class ClassPropertyBuilder { - setType(fieldType: FieldType): ClassPropertyBuilder - alias(alias?: string | undefined | null): ClassPropertyBuilder - description(description?: string | undefined | null): ClassPropertyBuilder -} - -export class EnumBuilder { - value(name: string): EnumValueBuilder - alias(alias?: string | undefined | null): EnumBuilder - field(): FieldType -} - -export class EnumValueBuilder { - alias(alias?: string | undefined | null): EnumValueBuilder - skip(skip?: boolean | undefined | null): EnumValueBuilder - description(description?: string | undefined | null): EnumValueBuilder -} - -export class FieldType { - list(): FieldType - optional(): FieldType -} - -export class FunctionResult { - parsed(): any -} - -export class FunctionResultStream { - onEvent(func: (err: any, param: FunctionResult) => void): void - done(rctx: RuntimeContextManager): Promise -} - -export class RuntimeContextManager { - upsertTags(tags: any): void - deepClone(): RuntimeContextManager -} - -export class TypeBuilder { - constructor() - getEnum(name: string): EnumBuilder - getClass(name: string): ClassBuilder - list(inner: FieldType): FieldType - optional(inner: FieldType): FieldType - string(): FieldType - int(): FieldType - float(): FieldType - bool(): FieldType - null(): FieldType -} - -export interface BamlLogEvent { - metadata: LogEventMetadata - prompt?: string - rawOutput?: string - parsedOutput?: string - startTime: string -} - -export function invoke_runtime_cli(params: Array): void - -export interface LogEventMetadata { - eventId: string - parentId?: string - rootEventId: string -} - diff --git a/engine/language_client_typescript/native.js b/engine/language_client_typescript/native.js deleted file mode 100644 index aae967a38..000000000 --- a/engine/language_client_typescript/native.js +++ /dev/null @@ -1,377 +0,0 @@ -// prettier-ignore -/* eslint-disable */ -/* auto-generated by NAPI-RS */ - -const { readFileSync } = require('fs') - -let nativeBinding = null -const loadErrors = [] - -const isMusl = () => { - let musl = false - if (process.platform === 'linux') { - musl = isMuslFromFilesystem() - if (musl === null) { - musl = isMuslFromReport() - } - if (musl === null) { - musl = isMuslFromChildProcess() - } - } - return musl -} - -const isFileMusl = (f) => f.includes('libc.musl-') || f.includes('ld-musl-') - -const isMuslFromFilesystem = () => { - try { - return readFileSync('/usr/bin/ldd', 'utf-8').includes('musl') - } catch { - return null - } -} - -const isMuslFromReport = () => { - const report = typeof process.report.getReport === 'function' ? process.report.getReport() : null - if (!report) { - return null - } - if (report.header && report.header.glibcVersionRuntime) { - return false - } - if (Array.isArray(report.sharedObjects)) { - if (report.sharedObjects.some(isFileMusl)) { - return true - } - } - return false -} - -const isMuslFromChildProcess = () => { - try { - return require('child_process').execSync('ldd --version', { encoding: 'utf8' }).includes('musl') - } catch (e) { - // If we reach this case, we don't know if the system is musl or not, so is better to just fallback to false - return false - } -} - -function requireNative() { - if (process.platform === 'android') { - if (process.arch === 'arm64') { - try { - return require('./baml.android-arm64.node') - } catch (e) { - loadErrors.push(e) - } - try { - return require('@boundaryml/baml-android-arm64') - } catch (e) { - loadErrors.push(e) - } - - } else if (process.arch === 'arm') { - try { - return require('./baml.android-arm-eabi.node') - } catch (e) { - loadErrors.push(e) - } - try { - return require('@boundaryml/baml-android-arm-eabi') - } catch (e) { - loadErrors.push(e) - } - - } else { - loadErrors.push(new Error(`Unsupported architecture on Android ${process.arch}`)) - } - } else if (process.platform === 'win32') { - if (process.arch === 'x64') { - try { - return require('./baml.win32-x64-msvc.node') - } catch (e) { - loadErrors.push(e) - } - try { - return require('@boundaryml/baml-win32-x64-msvc') - } catch (e) { - loadErrors.push(e) - } - - } else if (process.arch === 'ia32') { - try { - return require('./baml.win32-ia32-msvc.node') - } catch (e) { - loadErrors.push(e) - } - try { - return require('@boundaryml/baml-win32-ia32-msvc') - } catch (e) { - loadErrors.push(e) - } - - } else if (process.arch === 'arm64') { - try { - return require('./baml.win32-arm64-msvc.node') - } catch (e) { - loadErrors.push(e) - } - try { - return require('@boundaryml/baml-win32-arm64-msvc') - } catch (e) { - loadErrors.push(e) - } - - } else { - loadErrors.push(new Error(`Unsupported architecture on Windows: ${process.arch}`)) - } - } else if (process.platform === 'darwin') { - try { - return require('./baml.darwin-universal.node') - } catch (e) { - loadErrors.push(e) - } - try { - return require('@boundaryml/baml-darwin-universal') - } catch (e) { - loadErrors.push(e) - } - - if (process.arch === 'x64') { - try { - return require('./baml.darwin-x64.node') - } catch (e) { - loadErrors.push(e) - } - try { - return require('@boundaryml/baml-darwin-x64') - } catch (e) { - loadErrors.push(e) - } - - } else if (process.arch === 'arm64') { - try { - return require('./baml.darwin-arm64.node') - } catch (e) { - loadErrors.push(e) - } - try { - return require('@boundaryml/baml-darwin-arm64') - } catch (e) { - loadErrors.push(e) - } - - } else { - loadErrors.push(new Error(`Unsupported architecture on macOS: ${process.arch}`)) - } - } else if (process.platform === 'freebsd') { - if (process.arch === 'x64') { - try { - return require('./baml.freebsd-x64.node') - } catch (e) { - loadErrors.push(e) - } - try { - return require('@boundaryml/baml-freebsd-x64') - } catch (e) { - loadErrors.push(e) - } - - } else if (process.arch === 'arm64') { - try { - return require('./baml.freebsd-arm64.node') - } catch (e) { - loadErrors.push(e) - } - try { - return require('@boundaryml/baml-freebsd-arm64') - } catch (e) { - loadErrors.push(e) - } - - } else { - loadErrors.push(new Error(`Unsupported architecture on FreeBSD: ${process.arch}`)) - } - } else if (process.platform === 'linux') { - if (process.arch === 'x64') { - if (isMusl()) { - try { - return require('./baml.linux-x64-musl.node') - } catch (e) { - loadErrors.push(e) - } - try { - return require('@boundaryml/baml-linux-x64-musl') - } catch (e) { - loadErrors.push(e) - } - - } else { - try { - return require('./baml.linux-x64-gnu.node') - } catch (e) { - loadErrors.push(e) - } - try { - return require('@boundaryml/baml-linux-x64-gnu') - } catch (e) { - loadErrors.push(e) - } - - } - } else if (process.arch === 'arm64') { - if (isMusl()) { - try { - return require('./baml.linux-arm64-musl.node') - } catch (e) { - loadErrors.push(e) - } - try { - return require('@boundaryml/baml-linux-arm64-musl') - } catch (e) { - loadErrors.push(e) - } - - } else { - try { - return require('./baml.linux-arm64-gnu.node') - } catch (e) { - loadErrors.push(e) - } - try { - return require('@boundaryml/baml-linux-arm64-gnu') - } catch (e) { - loadErrors.push(e) - } - - } - } else if (process.arch === 'arm') { - if (isMusl()) { - try { - return require('./baml.linux-arm-musleabihf.node') - } catch (e) { - loadErrors.push(e) - } - try { - return require('@boundaryml/baml-linux-arm-musleabihf') - } catch (e) { - loadErrors.push(e) - } - - } else { - try { - return require('./baml.linux-arm-gnueabihf.node') - } catch (e) { - loadErrors.push(e) - } - try { - return require('@boundaryml/baml-linux-arm-gnueabihf') - } catch (e) { - loadErrors.push(e) - } - - } - } else if (process.arch === 'riscv64') { - if (isMusl()) { - try { - return require('./baml.linux-riscv64-musl.node') - } catch (e) { - loadErrors.push(e) - } - try { - return require('@boundaryml/baml-linux-riscv64-musl') - } catch (e) { - loadErrors.push(e) - } - - } else { - try { - return require('./baml.linux-riscv64-gnu.node') - } catch (e) { - loadErrors.push(e) - } - try { - return require('@boundaryml/baml-linux-riscv64-gnu') - } catch (e) { - loadErrors.push(e) - } - - } - } else if (process.arch === 'ppc64') { - try { - return require('./baml.linux-ppc64-gnu.node') - } catch (e) { - loadErrors.push(e) - } - try { - return require('@boundaryml/baml-linux-ppc64-gnu') - } catch (e) { - loadErrors.push(e) - } - - } else if (process.arch === 's390x') { - try { - return require('./baml.linux-s390x-gnu.node') - } catch (e) { - loadErrors.push(e) - } - try { - return require('@boundaryml/baml-linux-s390x-gnu') - } catch (e) { - loadErrors.push(e) - } - - } else { - loadErrors.push(new Error(`Unsupported architecture on Linux: ${process.arch}`)) - } - } else { - loadErrors.push(new Error(`Unsupported OS: ${process.platform}, architecture: ${process.arch}`)) - } -} - -nativeBinding = requireNative() - -if (!nativeBinding || process.env.NAPI_RS_FORCE_WASI) { - try { - nativeBinding = require('./baml.wasi.cjs') - } catch (err) { - if (process.env.NAPI_RS_FORCE_WASI) { - console.error(err) - } - } - if (!nativeBinding) { - try { - nativeBinding = require('@boundaryml/baml-wasm32-wasi') - } catch (err) { - if (process.env.NAPI_RS_FORCE_WASI) { - console.error(err) - } - } - } -} - -if (!nativeBinding) { - if (loadErrors.length > 0) { - // TODO Link to documentation with potential fixes - // - The package owner could build/publish bindings for this arch - // - The user may need to bundle the correct files - // - The user may need to re-install node_modules to get new packages - throw new Error('Failed to load native binding', { cause: loadErrors }) - } - throw new Error(`Failed to load native binding`) -} - -module.exports.BamlAudio = nativeBinding.BamlAudio -module.exports.BamlImage = nativeBinding.BamlImage -module.exports.BamlRuntime = nativeBinding.BamlRuntime -module.exports.BamlSpan = nativeBinding.BamlSpan -module.exports.ClassBuilder = nativeBinding.ClassBuilder -module.exports.ClassPropertyBuilder = nativeBinding.ClassPropertyBuilder -module.exports.EnumBuilder = nativeBinding.EnumBuilder -module.exports.EnumValueBuilder = nativeBinding.EnumValueBuilder -module.exports.FieldType = nativeBinding.FieldType -module.exports.FunctionResult = nativeBinding.FunctionResult -module.exports.FunctionResultStream = nativeBinding.FunctionResultStream -module.exports.RuntimeContextManager = nativeBinding.RuntimeContextManager -module.exports.TypeBuilder = nativeBinding.TypeBuilder -module.exports.invoke_runtime_cli = nativeBinding.invoke_runtime_cli diff --git a/engine/language_client_typescript/src/lib.rs b/engine/language_client_typescript/src/lib.rs index b5f2df72a..f7217a0a6 100644 --- a/engine/language_client_typescript/src/lib.rs +++ b/engine/language_client_typescript/src/lib.rs @@ -2,8 +2,11 @@ use napi::{Env, JsUndefined}; use napi_derive::napi; mod parse_ts_types; +mod runtime; mod types; +pub(crate) use runtime::BamlRuntime; + #[napi(js_name = "invoke_runtime_cli")] pub fn run_cli(env: Env, params: Vec) -> napi::Result { baml_runtime::BamlRuntime::run_cli(params, baml_runtime::CallerType::Typescript)?; diff --git a/engine/language_client_typescript/src/types/runtime.rs b/engine/language_client_typescript/src/runtime.rs similarity index 95% rename from engine/language_client_typescript/src/types/runtime.rs rename to engine/language_client_typescript/src/runtime.rs index 196848829..0148e9e31 100644 --- a/engine/language_client_typescript/src/types/runtime.rs +++ b/engine/language_client_typescript/src/runtime.rs @@ -1,9 +1,10 @@ -use super::function_result_stream::FunctionResultStream; -use super::runtime_ctx_manager::RuntimeContextManager; -use super::type_builder::TypeBuilder; use crate::parse_ts_types; +use crate::types::function_result_stream::FunctionResultStream; use crate::types::function_results::FunctionResult; -use baml_runtime::on_log_event::{LogEvent}; +use crate::types::runtime_ctx_manager::RuntimeContextManager; +use crate::types::trace_stats::TraceStats; +use crate::types::type_builder::TypeBuilder; +use baml_runtime::on_log_event::LogEvent; use baml_runtime::runtime_interface::ExperimentalTracingInterface; use baml_runtime::BamlRuntime as CoreRuntime; use baml_types::BamlValue; @@ -225,10 +226,11 @@ impl BamlRuntime { } #[napi] - pub fn flush(&mut self, env: Env) -> napi::Result<()> { + pub fn flush(&mut self, env: Env) -> napi::Result { let res = self .inner .flush() + .map(|stats| stats.into()) .map_err(|e| napi::Error::new(napi::Status::GenericFailure, e.to_string())); res diff --git a/engine/language_client_typescript/src/types/audio.rs b/engine/language_client_typescript/src/types/audio.rs index 87bbb1697..6a28c1489 100644 --- a/engine/language_client_typescript/src/types/audio.rs +++ b/engine/language_client_typescript/src/types/audio.rs @@ -2,6 +2,7 @@ use baml_types::BamlMediaType; use napi::bindgen_prelude::External; use napi_derive::napi; use serde_json::json; + crate::lang_wrapper!(BamlAudio, baml_types::BamlMedia); #[napi] diff --git a/engine/language_client_typescript/src/types/function_result_stream.rs b/engine/language_client_typescript/src/types/function_result_stream.rs index d12181c46..2fdf7fa29 100644 --- a/engine/language_client_typescript/src/types/function_result_stream.rs +++ b/engine/language_client_typescript/src/types/function_result_stream.rs @@ -18,7 +18,7 @@ crate::lang_wrapper!( ); impl FunctionResultStream { - pub(super) fn new( + pub(crate) fn new( inner: baml_runtime::FunctionResultStream, event: Option>, tb: Option, diff --git a/engine/language_client_typescript/src/types/mod.rs b/engine/language_client_typescript/src/types/mod.rs index 612e571f8..b5edf3562 100644 --- a/engine/language_client_typescript/src/types/mod.rs +++ b/engine/language_client_typescript/src/types/mod.rs @@ -1,10 +1,10 @@ mod lang_wrappers; pub mod audio; -mod function_result_stream; -mod function_results; +pub(crate) mod function_result_stream; +pub(crate) mod function_results; pub mod image; -mod runtime; -mod runtime_ctx_manager; -mod span; -mod type_builder; +pub(crate) mod runtime_ctx_manager; +pub(crate) mod span; +pub(crate) mod trace_stats; +pub(crate) mod type_builder; diff --git a/engine/language_client_typescript/src/types/span.rs b/engine/language_client_typescript/src/types/span.rs index ff7b526cd..56e89142b 100644 --- a/engine/language_client_typescript/src/types/span.rs +++ b/engine/language_client_typescript/src/types/span.rs @@ -2,8 +2,8 @@ use baml_runtime::runtime_interface::ExperimentalTracingInterface; use baml_types::BamlValue; use napi_derive::napi; -use super::runtime::BamlRuntime; use super::runtime_ctx_manager::RuntimeContextManager; +use crate::BamlRuntime; crate::lang_wrapper!(BamlSpan, Option>, diff --git a/engine/language_client_typescript/src/types/trace_stats.rs b/engine/language_client_typescript/src/types/trace_stats.rs new file mode 100644 index 000000000..7bcb47b1d --- /dev/null +++ b/engine/language_client_typescript/src/types/trace_stats.rs @@ -0,0 +1,11 @@ +use napi_derive::napi; + +crate::lang_wrapper!(TraceStats, baml_runtime::TraceStats); + +#[napi] +impl TraceStats { + #[napi] + pub fn n_spans_failed_before_submit(&self) -> u32 { + self.inner.n_spans_failed_before_submit + } +} diff --git a/engine/language_client_typescript/stream.d.ts b/engine/language_client_typescript/stream.d.ts deleted file mode 100644 index dd7f000cd..000000000 --- a/engine/language_client_typescript/stream.d.ts +++ /dev/null @@ -1,15 +0,0 @@ -import { FunctionResult, FunctionResultStream, RuntimeContextManager } from './native'; -export declare class BamlStream { - private ffiStream; - private partialCoerce; - private finalCoerce; - private ctxManager; - private task; - private eventQueue; - constructor(ffiStream: FunctionResultStream, partialCoerce: (result: FunctionResult) => PartialOutputType, finalCoerce: (result: FunctionResult) => FinalOutputType, ctxManager: RuntimeContextManager); - private driveToCompletion; - private driveToCompletionInBg; - [Symbol.asyncIterator](): AsyncIterableIterator; - getFinalResponse(): Promise; -} -//# sourceMappingURL=stream.d.ts.map \ No newline at end of file diff --git a/engine/language_client_typescript/stream.d.ts.map b/engine/language_client_typescript/stream.d.ts.map deleted file mode 100644 index 7c571df13..000000000 --- a/engine/language_client_typescript/stream.d.ts.map +++ /dev/null @@ -1 +0,0 @@ -{"version":3,"file":"stream.d.ts","sourceRoot":"","sources":["typescript_src/stream.ts"],"names":[],"mappings":"AAAA,OAAO,EAAE,cAAc,EAAE,oBAAoB,EAAE,qBAAqB,EAAE,MAAM,UAAU,CAAA;AAEtF,qBAAa,UAAU,CAAC,iBAAiB,EAAE,eAAe;IAMtD,OAAO,CAAC,SAAS;IACjB,OAAO,CAAC,aAAa;IACrB,OAAO,CAAC,WAAW;IACnB,OAAO,CAAC,UAAU;IARpB,OAAO,CAAC,IAAI,CAAuC;IAEnD,OAAO,CAAC,UAAU,CAAgC;gBAGxC,SAAS,EAAE,oBAAoB,EAC/B,aAAa,EAAE,CAAC,MAAM,EAAE,cAAc,KAAK,iBAAiB,EAC5D,WAAW,EAAE,CAAC,MAAM,EAAE,cAAc,KAAK,eAAe,EACxD,UAAU,EAAE,qBAAqB;YAG7B,iBAAiB;IAiB/B,OAAO,CAAC,qBAAqB;IAQtB,CAAC,MAAM,CAAC,aAAa,CAAC,IAAI,qBAAqB,CAAC,iBAAiB,CAAC;IAmBnE,gBAAgB,IAAI,OAAO,CAAC,eAAe,CAAC;CAKnD"} \ No newline at end of file diff --git a/engine/language_client_typescript/stream.js b/engine/language_client_typescript/stream.js deleted file mode 100644 index b7d6ae5f1..000000000 --- a/engine/language_client_typescript/stream.js +++ /dev/null @@ -1,59 +0,0 @@ -"use strict"; -Object.defineProperty(exports, "__esModule", { value: true }); -exports.BamlStream = void 0; -class BamlStream { - ffiStream; - partialCoerce; - finalCoerce; - ctxManager; - task = null; - eventQueue = []; - constructor(ffiStream, partialCoerce, finalCoerce, ctxManager) { - this.ffiStream = ffiStream; - this.partialCoerce = partialCoerce; - this.finalCoerce = finalCoerce; - this.ctxManager = ctxManager; - } - async driveToCompletion() { - try { - this.ffiStream.onEvent((err, data) => { - if (err) { - return; - } - else { - this.eventQueue.push(data); - } - }); - const retval = await this.ffiStream.done(this.ctxManager); - return retval; - } - finally { - this.eventQueue.push(null); - } - } - driveToCompletionInBg() { - if (this.task === null) { - this.task = this.driveToCompletion(); - } - return this.task; - } - async *[Symbol.asyncIterator]() { - this.driveToCompletionInBg(); - while (true) { - const event = this.eventQueue.shift(); - if (event === undefined) { - await new Promise((resolve) => setTimeout(resolve, 100)); - continue; - } - if (event === null) { - break; - } - yield this.partialCoerce(event.parsed()); - } - } - async getFinalResponse() { - const final = await this.driveToCompletionInBg(); - return this.finalCoerce(final.parsed()); - } -} -exports.BamlStream = BamlStream; diff --git a/engine/language_client_typescript/type_builder.d.ts b/engine/language_client_typescript/type_builder.d.ts deleted file mode 100644 index 9a9c70b81..000000000 --- a/engine/language_client_typescript/type_builder.d.ts +++ /dev/null @@ -1,51 +0,0 @@ -import { ClassPropertyBuilder as _ClassPropertyBuilder, EnumValueBuilder, FieldType, TypeBuilder as _TypeBuilder } from './native'; -type IsLiteral = string extends T ? false : true; -type NameOf = IsLiteral extends true ? T : 'DynamicType'; -type CheckNever = [T] extends [never] ? `Error: Attempt to add value '${Value}' which is already a part of '${NameOf}'.` : T; -type ExcludeFrom = T extends U ? never : T; -type RestrictNot = IsLiteral extends true ? CheckNever, Name, Value> : Value; -export declare class TypeBuilder { - private tb; - protected classes: Set; - protected enums: Set; - constructor({ classes, enums }: { - classes: Set; - enums: Set; - }); - _tb(): _TypeBuilder; - string(): FieldType; - int(): FieldType; - float(): FieldType; - bool(): FieldType; - list(type: FieldType): FieldType; - classBuilder(name: Name, properties: Properties[]): ClassBuilder; - enumBuilder(name: Name, values: T[]): EnumBuilder; - addClass(name: Name): ClassBuilder; - addEnum(name: Name): EnumBuilder; -} -export declare class ClassBuilder { - private properties; - private bldr; - constructor(tb: _TypeBuilder, name: ClassName, properties?: Set); - type(): FieldType; - listProperties(): Array<[string, ClassPropertyBuilder]>; - addProperty(name: RestrictNot, type: FieldType): ClassPropertyBuilder; - property(name: string): ClassPropertyBuilder; -} -declare class ClassPropertyBuilder { - private bldr; - constructor(bldr: _ClassPropertyBuilder); - alias(alias: string | null): ClassPropertyBuilder; - description(description: string | null): ClassPropertyBuilder; -} -export declare class EnumBuilder { - private values; - private bldr; - constructor(tb: _TypeBuilder, name: EnumName, values?: Set); - type(): FieldType; - value(name: S | T): EnumValueBuilder; - listValues(): Array<[string, EnumValueBuilder]>; - addValue(name: RestrictNot): EnumValueBuilder; -} -export {}; -//# sourceMappingURL=type_builder.d.ts.map \ No newline at end of file diff --git a/engine/language_client_typescript/type_builder.d.ts.map b/engine/language_client_typescript/type_builder.d.ts.map deleted file mode 100644 index c26a52d3a..000000000 --- a/engine/language_client_typescript/type_builder.d.ts.map +++ /dev/null @@ -1 +0,0 @@ -{"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 diff --git a/engine/language_client_typescript/type_builder.js b/engine/language_client_typescript/type_builder.js deleted file mode 100644 index bde0e4e18..000000000 --- a/engine/language_client_typescript/type_builder.js +++ /dev/null @@ -1,129 +0,0 @@ -"use strict"; -Object.defineProperty(exports, "__esModule", { value: true }); -exports.EnumBuilder = exports.ClassBuilder = exports.TypeBuilder = void 0; -const native_1 = require("./native"); -class TypeBuilder { - tb; - classes; - enums; - constructor({ classes, enums }) { - this.classes = classes; - this.enums = enums; - this.tb = new native_1.TypeBuilder(); - } - _tb() { - return this.tb; - } - string() { - return this.tb.string(); - } - int() { - return this.tb.int(); - } - float() { - return this.tb.float(); - } - bool() { - return this.tb.bool(); - } - list(type) { - return this.tb.list(type); - } - classBuilder(name, properties) { - return new ClassBuilder(this.tb, name, new Set(properties)); - } - enumBuilder(name, values) { - return new EnumBuilder(this.tb, name, new Set(values)); - } - addClass(name) { - if (this.classes.has(name)) { - throw new Error(`Class ${name} already exists`); - } - if (this.enums.has(name)) { - throw new Error(`Enum ${name} already exists`); - } - this.classes.add(name); - return new ClassBuilder(this.tb, name); - } - addEnum(name) { - if (this.classes.has(name)) { - throw new Error(`Class ${name} already exists`); - } - if (this.enums.has(name)) { - throw new Error(`Enum ${name} already exists`); - } - this.enums.add(name); - return new EnumBuilder(this.tb, name); - } -} -exports.TypeBuilder = TypeBuilder; -class ClassBuilder { - properties; - bldr; - constructor(tb, name, properties = new Set()) { - this.properties = properties; - this.bldr = tb.getClass(name); - } - type() { - return this.bldr.field(); - } - listProperties() { - return Array.from(this.properties).map((name) => [name, new ClassPropertyBuilder(this.bldr.property(name))]); - } - addProperty(name, type) { - if (this.properties.has(name)) { - throw new Error(`Property ${name} already exists.`); - } - this.properties.add(name); - return new ClassPropertyBuilder(this.bldr.property(name).setType(type)); - } - property(name) { - if (!this.properties.has(name)) { - throw new Error(`Property ${name} not found.`); - } - return new ClassPropertyBuilder(this.bldr.property(name)); - } -} -exports.ClassBuilder = ClassBuilder; -class ClassPropertyBuilder { - bldr; - constructor(bldr) { - this.bldr = bldr; - } - alias(alias) { - this.bldr.alias(alias); - return this; - } - description(description) { - this.bldr.description(description); - return this; - } -} -class EnumBuilder { - values; - bldr; - constructor(tb, name, values = new Set()) { - this.values = values; - this.bldr = tb.getEnum(name); - } - type() { - return this.bldr.field(); - } - value(name) { - if (!this.values.has(name)) { - throw new Error(`Value ${name} not found.`); - } - return this.bldr.value(name); - } - listValues() { - return Array.from(this.values).map((name) => [name, this.bldr.value(name)]); - } - addValue(name) { - if (this.values.has(name)) { - throw new Error(`Value ${name} already exists.`); - } - this.values.add(name); - return this.bldr.value(name); - } -} -exports.EnumBuilder = EnumBuilder; diff --git a/engine/language_client_typescript/typescript_src/async_context_vars.ts b/engine/language_client_typescript/typescript_src/async_context_vars.ts index a33f5c90c..2eb6ccdb4 100644 --- a/engine/language_client_typescript/typescript_src/async_context_vars.ts +++ b/engine/language_client_typescript/typescript_src/async_context_vars.ts @@ -1,7 +1,7 @@ import { BamlSpan, RuntimeContextManager, BamlRuntime, BamlLogEvent } from './native' import { AsyncLocalStorage } from 'async_hooks' -export class CtxManager { +export class BamlCtxManager { private rt: BamlRuntime private ctx: AsyncLocalStorage @@ -19,27 +19,18 @@ export class CtxManager { manager.upsertTags(tags) } - get(): RuntimeContextManager { + cloneContext(): RuntimeContextManager { let store = this.ctx.getStore() if (store === undefined) { store = this.rt.createContextManager() this.ctx.enterWith(store) } - return store + return store.deepClone() } - startTraceSync(name: string, args: Record): BamlSpan { - const mng = this.get() - // const clone = mng.deepClone() - // this.ctx.enterWith(clone) - return BamlSpan.new(this.rt, name, args, mng) - } - - startTraceAsync(name: string, args: Record): BamlSpan { - const mng = this.get() - const clone = mng.deepClone() - this.ctx.enterWith(clone) - return BamlSpan.new(this.rt, name, args, clone) + startTrace(name: string, args: Record): [RuntimeContextManager, BamlSpan] { + const mng = this.cloneContext() + return [mng, BamlSpan.new(this.rt, name, args, mng)] } endTrace(span: BamlSpan, response: any): void { @@ -48,7 +39,11 @@ export class CtxManager { console.error('Context lost before span could be finished\n') return } - span.finish(response, manager) + try { + span.finish(response, manager) + } catch (e) { + console.error('BAML internal error', e) + } } flush(): void { @@ -72,20 +67,21 @@ export class CtxManager { }), {}, ) - const span = this.startTraceSync(name, params) - - try { - const response = func(...args) - this.endTrace(span, response) - return response - } catch (e) { - this.endTrace(span, e) - throw e - } + const [mng, span] = this.startTrace(name, params) + this.ctx.run(mng, () => { + try { + const response = func(...args) + this.endTrace(span, response) + return response + } catch (e) { + this.endTrace(span, e) + throw e + } + }) }) } - traceFnAync Promise>(name: string, func: F): F { + traceFnAsync Promise>(name: string, func: F): F { const funcName = name return (async (...args: any[]) => { const params = args.reduce( @@ -95,15 +91,17 @@ export class CtxManager { }), {}, ) - const span = this.startTraceAsync(funcName, params) - try { - const response = await func(...args) - this.endTrace(span, response) - return response - } catch (e) { - this.endTrace(span, e) - throw e - } + const [mng, span] = this.startTrace(name, params) + await this.ctx.run(mng, async () => { + try { + const response = await func(...args) + this.endTrace(span, response) + return response + } catch (e) { + this.endTrace(span, e) + throw e + } + }) }) } } diff --git a/engine/language_client_typescript/typescript_src/index.ts b/engine/language_client_typescript/typescript_src/index.ts index 3ee94afce..6e19e66b6 100644 --- a/engine/language_client_typescript/typescript_src/index.ts +++ b/engine/language_client_typescript/typescript_src/index.ts @@ -1,3 +1,10 @@ -export { BamlRuntime, FunctionResult, FunctionResultStream, BamlImage as Image, BamlAudio as Audio, invoke_runtime_cli } from './native' +export { + BamlRuntime, + FunctionResult, + FunctionResultStream, + BamlImage as Image, + BamlAudio as Audio, + invoke_runtime_cli, +} from './native' export { BamlStream } from './stream' -export { CtxManager as BamlCtxManager } from './async_context_vars' +export { BamlCtxManager } from './async_context_vars' diff --git a/integ-tests/python/__init__.py b/integ-tests/python/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/integ-tests/python/baml_client/tracing.py b/integ-tests/python/baml_client/tracing.py index 3ac276994..7c2bff4d4 100644 --- a/integ-tests/python/baml_client/tracing.py +++ b/integ-tests/python/baml_client/tracing.py @@ -17,7 +17,8 @@ trace = DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_CTX.trace_fn set_tags = DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_CTX.upsert_tags -flush = DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_CTX.flush +def flush(): + DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_CTX.flush() on_log_event = DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_CTX.on_log_event diff --git a/integ-tests/python/poetry.lock b/integ-tests/python/poetry.lock index 8720de6ce..8850db6e1 100644 --- a/integ-tests/python/poetry.lock +++ b/integ-tests/python/poetry.lock @@ -14,6 +14,16 @@ files = [ [package.dependencies] typing-extensions = {version = ">=4.0.0", markers = "python_version < \"3.9\""} +[[package]] +name = "assertpy" +version = "1.1" +description = "Simple assertion library for unit testing in python with a fluent API" +optional = false +python-versions = "*" +files = [ + {file = "assertpy-1.1.tar.gz", hash = "sha256:acc64329934ad71a3221de185517a43af33e373bb44dc05b5a9b174394ef4833"}, +] + [[package]] name = "colorama" version = "0.4.6" @@ -320,4 +330,4 @@ files = [ [metadata] lock-version = "2.0" python-versions = "^3.8" -content-hash = "c0890bb7d42663b172b862caa5d317f325f92d8bd1e8e6a9e211461b189f22dc" +content-hash = "75bc9294db7074ba805acc12ba1a8a201c8099a40f8ffe1d2e2aefb14d48e5e6" diff --git a/integ-tests/python/pyproject.toml b/integ-tests/python/pyproject.toml index e50d88083..8c58d850b 100644 --- a/integ-tests/python/pyproject.toml +++ b/integ-tests/python/pyproject.toml @@ -19,6 +19,7 @@ pytest-asyncio = "^0.23.7" pytest = "^8.2.1" pydantic = "^2.7.1" python-dotenv = "^1.0.1" +assertpy = "^1.1" [build-system] requires = ["poetry-core"] diff --git a/integ-tests/python/base64_test_data.py b/integ-tests/python/tests/base64_test_data.py similarity index 100% rename from integ-tests/python/base64_test_data.py rename to integ-tests/python/tests/base64_test_data.py diff --git a/integ-tests/python/test_functions.py b/integ-tests/python/tests/test_functions.py similarity index 93% rename from integ-tests/python/test_functions.py rename to integ-tests/python/tests/test_functions.py index e33b13b23..bc355cb5d 100644 --- a/integ-tests/python/test_functions.py +++ b/integ-tests/python/tests/test_functions.py @@ -1,15 +1,16 @@ -from typing import List +import time import pytest +from assertpy import assert_that from dotenv import load_dotenv -from base64_test_data import image_b64, audio_b64 +from .base64_test_data import image_b64, audio_b64 load_dotenv() import baml_py -from baml_client import b -from baml_client.types import NamedArgsSingleEnumList, NamedArgsSingleClass -from baml_client.tracing import trace, set_tags, flush, on_log_event -from baml_client.type_builder import TypeBuilder +from ..baml_client import b +from ..baml_client.types import NamedArgsSingleEnumList, NamedArgsSingleClass +from ..baml_client.tracing import trace, set_tags, flush, on_log_event +from ..baml_client.type_builder import TypeBuilder import datetime @@ -180,12 +181,12 @@ async def test_streaming(): stream = b.stream.PromptTestStreaming( input="Programming languages are fun to create" ) - msgs = [] + msgs: list[str] = [] start_time = asyncio.get_event_loop().time() last_msg_time = start_time async for msg in stream: - msgs.append(msg) + msgs.append(str(msg)) if len(msgs) == 1: first_msg_time = asyncio.get_event_loop().time() @@ -222,9 +223,9 @@ async def test_streaming_uniterated(): @pytest.mark.asyncio async def test_streaming_claude(): stream = b.stream.PromptTestClaude(input="Mt Rainier is tall") - msgs = [] + msgs: list[str] = [] async for msg in stream: - msgs.append(msg) + msgs.append(str(msg)) final = await stream.get_final_response() assert len(final) > 0, "Expected non-empty final but got empty." @@ -246,7 +247,7 @@ async def test_streaming_claude(): @pytest.mark.asyncio async def test_streaming_gemini(): stream = b.stream.TestGemini(input="Dr.Pepper") - msgs: List[str] = [] + msgs: list[str] = [] async for msg in stream: if msg is not None: msgs.append(msg) @@ -270,11 +271,31 @@ async def test_streaming_gemini(): @pytest.mark.asyncio async def test_tracing_async(): - # sync_dummy_func("second-dummycall-arg") - res = await parent_async("first-arg-value") - # sync_dummy_func("second-dummycall-arg") - res2 = await parent_async2("second-arg-value") - # sync_dummy_func("second-dummycall-arg") + + @trace + async def nested_dummy_fn(_foo: str): + time.sleep(0.5 + random.random()) + return "nested dummy fn" + + @trace + async def dummy_fn(foo: str): + await asyncio.gather( + b.FnOutputClass(foo), + nested_dummy_fn(foo), + ) + return "dummy fn" + + await asyncio.gather( + dummy_fn("dummy arg 1"), + dummy_fn("dummy arg 2"), + dummy_fn("dummy arg 3"), + ) + await asyncio.gather( + parent_async("first-arg-value"), + parent_async2("second-arg-value") + ) + + assert_that(b._BamlClient__runtime.flush().n_spans_failed_before_submit()).is_equal_to(0) def test_tracing_sync(): diff --git a/integ-tests/typescript/baml_client/client.ts b/integ-tests/typescript/baml_client/client.ts index 6ca5545da..4d777d02b 100644 --- a/integ-tests/typescript/baml_client/client.ts +++ b/integ-tests/typescript/baml_client/client.ts @@ -26,9 +26,13 @@ export type RecursivePartialNull = T extends object : T | null; export class BamlClient { + private runtime: BamlRuntime + private ctx_manager: BamlCtxManager private stream_client: BamlStreamClient - constructor(private runtime: BamlRuntime, private ctx_manager: BamlCtxManager) { + constructor(runtime: BamlRuntime, ctx_manager: BamlCtxManager) { + this.runtime = runtime + this.ctx_manager = ctx_manager this.stream_client = new BamlStreamClient(runtime, ctx_manager) } @@ -46,7 +50,7 @@ export class BamlClient { { "aud": aud }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as string @@ -61,7 +65,7 @@ export class BamlClient { { "input": input }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as Category @@ -76,7 +80,7 @@ export class BamlClient { { "input": input }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as Category @@ -91,7 +95,7 @@ export class BamlClient { { "input": input }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as Category @@ -106,7 +110,7 @@ export class BamlClient { { "img": img }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as string @@ -121,7 +125,7 @@ export class BamlClient { { "classWithImage": classWithImage,"img2": img2 }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as string @@ -136,7 +140,7 @@ export class BamlClient { { "classWithImage": classWithImage,"img2": img2 }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as string @@ -151,7 +155,7 @@ export class BamlClient { { "classWithImage": classWithImage,"img2": img2 }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as string @@ -166,7 +170,7 @@ export class BamlClient { { "input": input }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as DynamicClassTwo @@ -181,7 +185,7 @@ export class BamlClient { { "input": input }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as DynInputOutput @@ -196,7 +200,7 @@ export class BamlClient { { "input": input }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as DynInputOutput[] @@ -211,7 +215,7 @@ export class BamlClient { { "input": input }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as string[] @@ -226,7 +230,7 @@ export class BamlClient { { "text": text }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as Person[] @@ -241,7 +245,7 @@ export class BamlClient { { "email": email }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as ReceiptInfo @@ -256,7 +260,7 @@ export class BamlClient { { "resume": resume,"img": img?? null }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as Resume @@ -271,7 +275,7 @@ export class BamlClient { { "resume": resume }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as Resume @@ -286,7 +290,7 @@ export class BamlClient { { "input": input }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as ClassOptionalOutput | null @@ -301,7 +305,7 @@ export class BamlClient { { "input": input }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as ClassOptionalOutput2 | null @@ -316,7 +320,7 @@ export class BamlClient { { "input": input }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as EnumOutput[] @@ -331,7 +335,7 @@ export class BamlClient { { "input": input }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as EnumOutput @@ -346,7 +350,7 @@ export class BamlClient { { "myString": myString?? null }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as string @@ -361,7 +365,7 @@ export class BamlClient { { "input": input }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as boolean @@ -376,7 +380,7 @@ export class BamlClient { { "input": input }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as TestOutputClass @@ -391,7 +395,7 @@ export class BamlClient { { "input": input }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as TestOutputClass[] @@ -406,7 +410,7 @@ export class BamlClient { { "input": input }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as TestClassNested @@ -421,7 +425,7 @@ export class BamlClient { { "input": input }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as TestClassWithEnum @@ -436,7 +440,7 @@ export class BamlClient { { "input": input }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as string[] @@ -451,7 +455,7 @@ export class BamlClient { { "input": input }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as TestEnum @@ -466,7 +470,7 @@ export class BamlClient { { "input": input }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as TestClassAlias @@ -481,7 +485,7 @@ export class BamlClient { { "myArg": myArg }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as string @@ -496,7 +500,7 @@ export class BamlClient { { "text": text }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as RaysData @@ -511,7 +515,7 @@ export class BamlClient { { "email": email }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as OrderInfo @@ -526,7 +530,7 @@ export class BamlClient { { "query": query }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as SearchParams @@ -541,7 +545,7 @@ export class BamlClient { { "input": input }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as DynamicOutput @@ -556,7 +560,7 @@ export class BamlClient { { "input": input }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as (OptionalTest_ReturnType | null)[] @@ -571,7 +575,7 @@ export class BamlClient { { "input": input }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as string @@ -586,7 +590,7 @@ export class BamlClient { { "input": input }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as string @@ -601,7 +605,7 @@ export class BamlClient { { "input": input }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as string @@ -616,7 +620,7 @@ export class BamlClient { { "input": input }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as string @@ -631,7 +635,7 @@ export class BamlClient { { "input": input }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as string @@ -646,7 +650,7 @@ export class BamlClient { { "input": input }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as string @@ -661,7 +665,7 @@ export class BamlClient { { "input": input }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as string @@ -676,7 +680,7 @@ export class BamlClient { { "input": input }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as string @@ -691,7 +695,7 @@ export class BamlClient { { "input": input }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as string @@ -706,7 +710,7 @@ export class BamlClient { { }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as string @@ -721,7 +725,7 @@ export class BamlClient { { "myBool": myBool }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as string @@ -736,7 +740,7 @@ export class BamlClient { { "myArg": myArg }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as string @@ -751,7 +755,7 @@ export class BamlClient { { "myArg": myArg }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as string @@ -766,7 +770,7 @@ export class BamlClient { { "myFloat": myFloat }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as string @@ -781,7 +785,7 @@ export class BamlClient { { "myInt": myInt }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as string @@ -796,7 +800,7 @@ export class BamlClient { { "myString": myString }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as string @@ -811,7 +815,7 @@ export class BamlClient { { "myStringArray": myStringArray }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as string @@ -826,7 +830,7 @@ export class BamlClient { { "myArg": myArg }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as string @@ -841,7 +845,7 @@ export class BamlClient { { "input": input }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as string @@ -856,7 +860,7 @@ export class BamlClient { { "img": img }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as string @@ -871,7 +875,7 @@ export class BamlClient { { "myArg": myArg,"myArg2": myArg2 }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as string @@ -886,7 +890,7 @@ export class BamlClient { { "input": input }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as string @@ -901,7 +905,7 @@ export class BamlClient { { "input": input }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as string @@ -916,7 +920,7 @@ export class BamlClient { { "input": input }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as string @@ -931,7 +935,7 @@ export class BamlClient { { }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as string @@ -946,7 +950,7 @@ export class BamlClient { { }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as string @@ -961,7 +965,7 @@ export class BamlClient { { "input": input }, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return raw.parsed() as UnionTest_ReturnType @@ -983,14 +987,14 @@ class BamlStreamClient { "aud": aud }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, string>( raw, (a): a is RecursivePartialNull => a, (a): a is string => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -1005,14 +1009,14 @@ class BamlStreamClient { "input": input }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, Category>( raw, (a): a is RecursivePartialNull => a, (a): a is Category => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -1027,14 +1031,14 @@ class BamlStreamClient { "input": input }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, Category>( raw, (a): a is RecursivePartialNull => a, (a): a is Category => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -1049,14 +1053,14 @@ class BamlStreamClient { "input": input }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, Category>( raw, (a): a is RecursivePartialNull => a, (a): a is Category => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -1071,14 +1075,14 @@ class BamlStreamClient { "img": img }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, string>( raw, (a): a is RecursivePartialNull => a, (a): a is string => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -1093,14 +1097,14 @@ class BamlStreamClient { "classWithImage": classWithImage,"img2": img2 }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, string>( raw, (a): a is RecursivePartialNull => a, (a): a is string => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -1115,14 +1119,14 @@ class BamlStreamClient { "classWithImage": classWithImage,"img2": img2 }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, string>( raw, (a): a is RecursivePartialNull => a, (a): a is string => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -1137,14 +1141,14 @@ class BamlStreamClient { "classWithImage": classWithImage,"img2": img2 }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, string>( raw, (a): a is RecursivePartialNull => a, (a): a is string => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -1159,14 +1163,14 @@ class BamlStreamClient { "input": input }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, DynamicClassTwo>( raw, (a): a is RecursivePartialNull => a, (a): a is DynamicClassTwo => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -1181,14 +1185,14 @@ class BamlStreamClient { "input": input }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, DynInputOutput>( raw, (a): a is RecursivePartialNull => a, (a): a is DynInputOutput => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -1203,14 +1207,14 @@ class BamlStreamClient { "input": input }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, DynInputOutput[]>( raw, (a): a is RecursivePartialNull => a, (a): a is DynInputOutput[] => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -1225,14 +1229,14 @@ class BamlStreamClient { "input": input }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, string[]>( raw, (a): a is RecursivePartialNull => a, (a): a is string[] => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -1247,14 +1251,14 @@ class BamlStreamClient { "text": text }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, Person[]>( raw, (a): a is RecursivePartialNull => a, (a): a is Person[] => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -1269,14 +1273,14 @@ class BamlStreamClient { "email": email }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, ReceiptInfo>( raw, (a): a is RecursivePartialNull => a, (a): a is ReceiptInfo => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -1291,14 +1295,14 @@ class BamlStreamClient { "resume": resume,"img": img ?? null }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, Resume>( raw, (a): a is RecursivePartialNull => a, (a): a is Resume => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -1313,14 +1317,14 @@ class BamlStreamClient { "resume": resume }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, Resume>( raw, (a): a is RecursivePartialNull => a, (a): a is Resume => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -1335,14 +1339,14 @@ class BamlStreamClient { "input": input }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, ClassOptionalOutput | null>( raw, (a): a is RecursivePartialNull => a, (a): a is ClassOptionalOutput | null => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -1357,14 +1361,14 @@ class BamlStreamClient { "input": input }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, ClassOptionalOutput2 | null>( raw, (a): a is RecursivePartialNull => a, (a): a is ClassOptionalOutput2 | null => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -1379,14 +1383,14 @@ class BamlStreamClient { "input": input }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, EnumOutput[]>( raw, (a): a is RecursivePartialNull => a, (a): a is EnumOutput[] => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -1401,14 +1405,14 @@ class BamlStreamClient { "input": input }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, EnumOutput>( raw, (a): a is RecursivePartialNull => a, (a): a is EnumOutput => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -1423,14 +1427,14 @@ class BamlStreamClient { "myString": myString ?? null }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, string>( raw, (a): a is RecursivePartialNull => a, (a): a is string => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -1445,14 +1449,14 @@ class BamlStreamClient { "input": input }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, boolean>( raw, (a): a is RecursivePartialNull => a, (a): a is boolean => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -1467,14 +1471,14 @@ class BamlStreamClient { "input": input }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, TestOutputClass>( raw, (a): a is RecursivePartialNull => a, (a): a is TestOutputClass => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -1489,14 +1493,14 @@ class BamlStreamClient { "input": input }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, TestOutputClass[]>( raw, (a): a is RecursivePartialNull => a, (a): a is TestOutputClass[] => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -1511,14 +1515,14 @@ class BamlStreamClient { "input": input }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, TestClassNested>( raw, (a): a is RecursivePartialNull => a, (a): a is TestClassNested => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -1533,14 +1537,14 @@ class BamlStreamClient { "input": input }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, TestClassWithEnum>( raw, (a): a is RecursivePartialNull => a, (a): a is TestClassWithEnum => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -1555,14 +1559,14 @@ class BamlStreamClient { "input": input }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, string[]>( raw, (a): a is RecursivePartialNull => a, (a): a is string[] => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -1577,14 +1581,14 @@ class BamlStreamClient { "input": input }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, TestEnum>( raw, (a): a is RecursivePartialNull => a, (a): a is TestEnum => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -1599,14 +1603,14 @@ class BamlStreamClient { "input": input }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, TestClassAlias>( raw, (a): a is RecursivePartialNull => a, (a): a is TestClassAlias => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -1621,14 +1625,14 @@ class BamlStreamClient { "myArg": myArg }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, string>( raw, (a): a is RecursivePartialNull => a, (a): a is string => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -1643,14 +1647,14 @@ class BamlStreamClient { "text": text }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, RaysData>( raw, (a): a is RecursivePartialNull => a, (a): a is RaysData => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -1665,14 +1669,14 @@ class BamlStreamClient { "email": email }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, OrderInfo>( raw, (a): a is RecursivePartialNull => a, (a): a is OrderInfo => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -1687,14 +1691,14 @@ class BamlStreamClient { "query": query }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, SearchParams>( raw, (a): a is RecursivePartialNull => a, (a): a is SearchParams => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -1709,14 +1713,14 @@ class BamlStreamClient { "input": input }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, DynamicOutput>( raw, (a): a is RecursivePartialNull => a, (a): a is DynamicOutput => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -1731,14 +1735,14 @@ class BamlStreamClient { "input": input }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, (OptionalTest_ReturnType | null)[]>( raw, (a): a is RecursivePartialNull<(OptionalTest_ReturnType | null)[]> => a, (a): a is (OptionalTest_ReturnType | null)[] => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -1753,14 +1757,14 @@ class BamlStreamClient { "input": input }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, string>( raw, (a): a is RecursivePartialNull => a, (a): a is string => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -1775,14 +1779,14 @@ class BamlStreamClient { "input": input }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, string>( raw, (a): a is RecursivePartialNull => a, (a): a is string => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -1797,14 +1801,14 @@ class BamlStreamClient { "input": input }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, string>( raw, (a): a is RecursivePartialNull => a, (a): a is string => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -1819,14 +1823,14 @@ class BamlStreamClient { "input": input }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, string>( raw, (a): a is RecursivePartialNull => a, (a): a is string => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -1841,14 +1845,14 @@ class BamlStreamClient { "input": input }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, string>( raw, (a): a is RecursivePartialNull => a, (a): a is string => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -1863,14 +1867,14 @@ class BamlStreamClient { "input": input }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, string>( raw, (a): a is RecursivePartialNull => a, (a): a is string => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -1885,14 +1889,14 @@ class BamlStreamClient { "input": input }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, string>( raw, (a): a is RecursivePartialNull => a, (a): a is string => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -1907,14 +1911,14 @@ class BamlStreamClient { "input": input }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, string>( raw, (a): a is RecursivePartialNull => a, (a): a is string => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -1929,14 +1933,14 @@ class BamlStreamClient { "input": input }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, string>( raw, (a): a is RecursivePartialNull => a, (a): a is string => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -1951,14 +1955,14 @@ class BamlStreamClient { }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, string>( raw, (a): a is RecursivePartialNull => a, (a): a is string => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -1973,14 +1977,14 @@ class BamlStreamClient { "myBool": myBool }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, string>( raw, (a): a is RecursivePartialNull => a, (a): a is string => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -1995,14 +1999,14 @@ class BamlStreamClient { "myArg": myArg }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, string>( raw, (a): a is RecursivePartialNull => a, (a): a is string => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -2017,14 +2021,14 @@ class BamlStreamClient { "myArg": myArg }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, string>( raw, (a): a is RecursivePartialNull => a, (a): a is string => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -2039,14 +2043,14 @@ class BamlStreamClient { "myFloat": myFloat }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, string>( raw, (a): a is RecursivePartialNull => a, (a): a is string => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -2061,14 +2065,14 @@ class BamlStreamClient { "myInt": myInt }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, string>( raw, (a): a is RecursivePartialNull => a, (a): a is string => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -2083,14 +2087,14 @@ class BamlStreamClient { "myString": myString }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, string>( raw, (a): a is RecursivePartialNull => a, (a): a is string => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -2105,14 +2109,14 @@ class BamlStreamClient { "myStringArray": myStringArray }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, string>( raw, (a): a is RecursivePartialNull => a, (a): a is string => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -2127,14 +2131,14 @@ class BamlStreamClient { "myArg": myArg }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, string>( raw, (a): a is RecursivePartialNull => a, (a): a is string => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -2149,14 +2153,14 @@ class BamlStreamClient { "input": input }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, string>( raw, (a): a is RecursivePartialNull => a, (a): a is string => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -2171,14 +2175,14 @@ class BamlStreamClient { "img": img }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, string>( raw, (a): a is RecursivePartialNull => a, (a): a is string => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -2193,14 +2197,14 @@ class BamlStreamClient { "myArg": myArg,"myArg2": myArg2 }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, string>( raw, (a): a is RecursivePartialNull => a, (a): a is string => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -2215,14 +2219,14 @@ class BamlStreamClient { "input": input }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, string>( raw, (a): a is RecursivePartialNull => a, (a): a is string => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -2237,14 +2241,14 @@ class BamlStreamClient { "input": input }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, string>( raw, (a): a is RecursivePartialNull => a, (a): a is string => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -2259,14 +2263,14 @@ class BamlStreamClient { "input": input }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, string>( raw, (a): a is RecursivePartialNull => a, (a): a is string => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -2281,14 +2285,14 @@ class BamlStreamClient { }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, string>( raw, (a): a is RecursivePartialNull => a, (a): a is string => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -2303,14 +2307,14 @@ class BamlStreamClient { }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, string>( raw, (a): a is RecursivePartialNull => a, (a): a is string => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } @@ -2325,14 +2329,14 @@ class BamlStreamClient { "input": input }, undefined, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) return new BamlStream, UnionTest_ReturnType>( raw, (a): a is RecursivePartialNull => a, (a): a is UnionTest_ReturnType => a, - this.ctx_manager.get(), + this.ctx_manager.cloneContext(), __baml_options__?.tb?.__tb(), ) } diff --git a/integ-tests/typescript/baml_client/tracing.ts b/integ-tests/typescript/baml_client/tracing.ts index 0a3fc3612..db3cc2deb 100644 --- a/integ-tests/typescript/baml_client/tracing.ts +++ b/integ-tests/typescript/baml_client/tracing.ts @@ -19,13 +19,14 @@ import { BamlLogEvent } from '@boundaryml/baml'; import { DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_CTX } from './globals'; const traceAsync = -DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_CTX.traceFnAync.bind(DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_CTX) +DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_CTX.traceFnAsync.bind(DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_CTX) const traceSync = DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_CTX.traceFnSync.bind(DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_CTX) const setTags = DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_CTX.upsertTags.bind(DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_CTX) -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 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) => DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_CTX.onLogEvent(callback) diff --git a/integ-tests/typescript/tests/integ-tests.test.ts b/integ-tests/typescript/tests/integ-tests.test.ts index 5cf5b6173..b023f7972 100644 --- a/integ-tests/typescript/tests/integ-tests.test.ts +++ b/integ-tests/typescript/tests/integ-tests.test.ts @@ -1,4 +1,5 @@ import assert from 'assert' +import { scheduler } from 'node:timers/promises' import { image_b64, audio_b64 } from './base64_test_data' import { Image } from '@boundaryml/baml' import { Audio } from '@boundaryml/baml' @@ -15,7 +16,8 @@ import { import TypeBuilder from '../baml_client/type_builder' import { RecursivePartialNull } from '../baml_client/client' import { config } from 'dotenv' -import { BamlLogEvent } from '@boundaryml/baml/native' +import { BamlLogEvent, BamlRuntime } from '@boundaryml/baml/native' +import { AsyncLocalStorage } from 'async_hooks' config() describe('Integ tests', () => { @@ -116,12 +118,12 @@ describe('Integ tests', () => { let res = await b.TestImageInput( Image.fromUrl('https://upload.wikimedia.org/wikipedia/en/4/4d/Shrek_%28character%29.png'), ) - expect(res.toLowerCase()).toContain('green') + expect(res.toLowerCase()).toMatch(/(green|yellow|ogre|shrek)/) }) it('should work with image from base 64', async () => { let res = await b.TestImageInput(Image.fromBase64('image/png', image_b64)) - expect(res.toLowerCase()).toContain('green') + expect(res.toLowerCase()).toMatch(/(green|yellow|ogre|shrek)/) }) it('should work with audio base 64', async () => { @@ -179,7 +181,7 @@ describe('Integ tests', () => { expect(msgs[i + 1].startsWith(msgs[i])).toBeTruthy() } expect(msgs.at(-1)).toEqual(final) - }) + }, 10_000) it('should support AWS', async () => { const res = await b.TestAws('Dr. Pepper') @@ -226,6 +228,8 @@ describe('Integ tests', () => { it('supports tracing sync', async () => { const blah = 'blah' + const dummyFunc = (_myArg: string): string => 'hello world' + const res = traceSync('myFuncParent', (firstArg: string, secondArg: number) => { setTags({ myKey: 'myVal' }) @@ -248,17 +252,44 @@ describe('Integ tests', () => { // Look at the dashboard to verify results. it('supports tracing async', async () => { + const nestedDummyFn = async (myArg: string): Promise => { + await scheduler.wait(100) // load-bearing: this ensures that we actually test concurrent execution + console.log('samDummyNested', myArg) + return myArg + } + + const dummyFn = async (myArg: string): Promise => { + await scheduler.wait(100) // load-bearing: this ensures that we actually test concurrent execution + const nested = await Promise.all([ + traceAsync('trace:nestedDummyFn1', nestedDummyFn)('nested1'), + traceAsync('trace:nestedDummyFn2', nestedDummyFn)('nested2'), + traceAsync('trace:nestedDummyFn3', nestedDummyFn)('nested3'), + ]) + console.log('dummy', myArg) + return myArg + } + + await Promise.all([ + traceAsync('trace:dummyFn1', dummyFn)('hi1'), + traceAsync('trace:dummyFn2', dummyFn)('hi2'), + traceAsync('trace:dummyFn3', dummyFn)('hi3'), + ]) + const res = await traceAsync('parentAsync', async (firstArg: string, secondArg: number) => { console.log('hello world') setTags({ myKey: 'myVal' }) - const res1 = traceSync('dummyFunc', dummyFunc)('firstDummyFuncArg') + const res1 = traceSync('dummyFunc', dummyFn)('firstDummyFuncArg') - const res2 = await traceAsync('asyncDummyFunc', asyncDummyFunc)('secondDummyFuncArg') + const res2 = await traceAsync('asyncDummyFunc', dummyFn)('secondDummyFuncArg') - const llm_res = await b.TestFnNamedArgsSingleStringList(['a', 'b', 'c']) + const llm_res = await Promise.all([ + b.TestFnNamedArgsSingleStringList(['a1', 'b', 'c']), + b.TestFnNamedArgsSingleStringList(['a2', 'b', 'c']), + b.TestFnNamedArgsSingleStringList(['a3', 'b', 'c']), + ]) - const res3 = await traceAsync('asyncDummyFunc', asyncDummyFunc)('thirdDummyFuncArg') + const res3 = await traceAsync('asyncDummyFunc', dummyFn)('thirdDummyFuncArg') return 'hello world' })('hi', 10) @@ -266,10 +297,14 @@ describe('Integ tests', () => { const res2 = await traceAsync('parentAsync2', async (firstArg: string, secondArg: number) => { console.log('hello world') - const res1 = traceSync('dummyFunc', dummyFunc)('firstDummyFuncArg') + const syncDummyFn = (_myArg: string): string => 'hello world' + const res1 = traceSync('dummyFunc', syncDummyFn)('firstDummyFuncArg') return 'hello world' })('hi', 10) + + const failedToSubmitCount = ((b as any).runtime as BamlRuntime).flush().nSpansFailedBeforeSubmit() + expect(failedToSubmitCount).toEqual(0) }) it('should work with dynamic types single', async () => { @@ -364,27 +399,12 @@ describe('Integ tests', () => { }) }) -function asyncDummyFunc(myArg: string): Promise { - console.log('asyncDummyFuncArgs', arguments) - return new Promise((resolve) => { - resolve({ - key: 'key', - key_two: true, - key_three: 52, - }) - }) -} - interface MyInterface { key: string key_two: boolean key_three: number } -function dummyFunc(myArg: string): string { - return 'hello world' -} - afterAll(async () => { flush() }) diff --git a/root.code-workspace b/root.code-workspace index 7c7b39cd1..1c3ac40a5 100644 --- a/root.code-workspace +++ b/root.code-workspace @@ -43,7 +43,9 @@ "editor.defaultFormatter": "biomejs.biome" }, "files.associations": { - "*.baml.j2": "jinja" + "*.baml.j2": "jinja", + "*.js.j2": "jinja-js", + "*.ts.j2": "jinja-js" }, "editor.colorDecoratorsLimit": 2000, "editor.formatOnSaveMode": "file", diff --git a/tools/build b/tools/build index cc8a00295..0652f4c7e 100755 --- a/tools/build +++ b/tools/build @@ -207,7 +207,7 @@ case "$_path" in command="env -u CONDA_PREFIX poetry run maturin develop --manifest-path ${_repo_root}/engine/language_client_python/Cargo.toml" command="${command} && poetry run baml-cli generate --from ${_repo_root}/integ-tests/baml_src" if [ "$_test_mode" -eq 1 ]; then - command="${command} && infisical run --env=test -- poetry run pytest" + command="${command} && infisical run --env=test -- poetry run pytest -k test_tracing_async" fi if [ "$_watch_mode" -eq 1 ]; then npx nodemon \ @@ -223,17 +223,26 @@ case "$_path" in ;; /integ-tests/typescript | /integ-tests/typescript/* ) + #BAML_LOG="baml_runtime::tracing=trace,baml_runtime::types::context_manager=debug,baml_runtime=debug" + BAML_LOG="baml_runtime=debug" command="(cd ${_repo_root}/engine/language_client_typescript && pnpm build:debug)" command="${command} && pnpm baml-cli generate --from ${_repo_root}/integ-tests/baml_src" if [ "$_test_mode" -eq 1 ]; then - command="${command} && pnpm integ-tests" + #command="${command} && pnpm integ-tests" + command="${command} && BAML_LOG=${BAML_LOG} infisical run -- pnpm test tests/integ-tests.test.ts" fi if [ "$_watch_mode" -eq 1 ]; then + #--verbose \ npx nodemon \ - --ext py,pyi,rs,j2,toml \ + --delay 1.5 \ + --ext py,pyi,rs,j2,toml,test.ts \ --watch "${_repo_root}/engine" \ --watch . \ - --ignore baml_client \ + --ignore baml_client/** \ + --ignore dist/** \ + --ignore target/** \ + --ignore node_modules/** \ + --ignore *.d.ts \ --exec "${command}" else eval "${command}"