try to fix typescript tracing

engine/baml-runtime/src/tracing/mod.rs

@@ -1,12 +1,9 @@
 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 baml_types::{BamlMap, BamlMediaType, BamlValue};
+use cfg_if::cfg_if;
 use colored::Colorize;
 use internal_baml_jinja::RenderedPrompt;
 use std::collections::HashMap;
@@ -27,16 +24,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,
@@ -95,6 +93,7 @@ impl BamlTracer {
         params: &BamlMap<String, BamlValue>,
     ) -> (Option<TracingSpan>, RuntimeContext) {
         let span_id = ctx.enter(function_name);
+        log::trace!("Entering span: {:#?}:::{:?}", span_id, function_name);
         if !self.enabled {
             return (None, ctx.create_ctx(tb));
         }
@@ -150,6 +149,12 @@ impl BamlTracer {
                 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);
@@ -211,6 +216,13 @@ impl BamlTracer {
             anyhow::bail!("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);
         }

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<TypescriptEnum<'ir>>,
     classes: Vec<TypescriptClass<'ir>>,
 }
 
 #[derive(askama::Template)]
-#[template(path = "types.js.j2", escape = "none")]
+#[template(path = "types.ts.j2", escape = "none")]
 pub(crate) struct TypescriptTypes<'ir> {
     enums: Vec<TypescriptEnum<'ir>>,
     classes: Vec<TypescriptClass<'ir>>,

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<TypescriptFunction>,
     types: Vec<String>,

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

@@ -11,9 +11,13 @@ export type RecursivePartialNull<T> = 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)
   }
 

engine/language_client_typescript/async_context_vars.d.ts

@@ -1,12 +1,12 @@
 import { BamlSpan, RuntimeContextManager, BamlRuntime, BamlLogEvent } from './native';
-export declare class CtxManager {
+export declare class BamlCtxManager {
     private rt;
     private ctx;
     constructor(rt: BamlRuntime);
     upsertTags(tags: Record<string, string>): void;
     get(): RuntimeContextManager;
-    startTraceSync(name: string, args: Record<string, any>): BamlSpan;
-    startTraceAsync(name: string, args: Record<string, any>): BamlSpan;
+    startTraceSync(name: string, args: Record<string, any>): [RuntimeContextManager, BamlSpan];
+    startTraceAsync(name: string, args: Record<string, any>): [RuntimeContextManager, BamlSpan];
     endTrace(span: BamlSpan, response: any): void;
     flush(): void;
     onLogEvent(callback: (event: BamlLogEvent) => void): void;

engine/language_client_typescript/async_context_vars.js

@@ -1,9 +1,9 @@
 "use strict";
 Object.defineProperty(exports, "__esModule", { value: true });
-exports.CtxManager = void 0;
+exports.BamlCtxManager = void 0;
 const native_1 = require("./native");
 const async_hooks_1 = require("async_hooks");
-class CtxManager {
+class BamlCtxManager {
     rt;
     ctx;
     constructor(rt) {
@@ -28,15 +28,11 @@ class CtxManager {
     }
     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);
+        return [mng, 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);
+        const mng = this.get().deepClone();
+        return [mng, native_1.BamlSpan.new(this.rt, name, args, mng)];
     }
     endTrace(span, response) {
         const manager = this.ctx.getStore();
@@ -62,16 +58,18 @@ class CtxManager {
                 ...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;
-            }
+            const [mng, span] = this.startTraceSync(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(name, func) {
@@ -81,17 +79,19 @@ class CtxManager {
                 ...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;
-            }
+            const [mng, span] = this.startTraceAsync(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;
+                }
+            });
         });
     }
 }
-exports.CtxManager = CtxManager;
+exports.BamlCtxManager = BamlCtxManager;

engine/language_client_typescript/index.d.ts

@@ -1,4 +1,4 @@
-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';
 //# sourceMappingURL=index.d.ts.map

engine/language_client_typescript/index.js

@@ -11,4 +11,4 @@ Object.defineProperty(exports, "invoke_runtime_cli", { enumerable: true, get: fu
 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; } });
+Object.defineProperty(exports, "BamlCtxManager", { enumerable: true, get: function () { return async_context_vars_1.BamlCtxManager; } });

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<RuntimeContextManager>
 
@@ -28,18 +28,14 @@ export class CtxManager {
     return store
   }
 
-  startTraceSync(name: string, args: Record<string, any>): BamlSpan {
+  startTraceSync(name: string, args: Record<string, any>): [RuntimeContextManager, BamlSpan] {
     const mng = this.get()
-    // const clone = mng.deepClone()
-    // this.ctx.enterWith(clone)
-    return BamlSpan.new(this.rt, name, args, mng)
+    return [mng, BamlSpan.new(this.rt, name, args, mng)]
   }
 
-  startTraceAsync(name: string, args: Record<string, any>): BamlSpan {
-    const mng = this.get()
-    const clone = mng.deepClone()
-    this.ctx.enterWith(clone)
-    return BamlSpan.new(this.rt, name, args, clone)
+  startTraceAsync(name: string, args: Record<string, any>): [RuntimeContextManager, BamlSpan] {
+    const mng = this.get().deepClone()
+    return [mng, BamlSpan.new(this.rt, name, args, mng)]
   }
 
   endTrace(span: BamlSpan, response: any): void {
@@ -72,16 +68,17 @@ 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.startTraceSync(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
+        }
+      })
     })
   }
 
@@ -95,15 +92,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.startTraceAsync(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
+        }
+      })
     })
   }
 }

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'

integ-tests/baml_src/clients.baml

@@ -84,18 +84,27 @@ client<llm> Gemini {
 }
 
 client<llm> AwsBedrock {
-  provider aws-bedrock
+  provider anthropic
   options {
-    inference_configuration {
-      max_tokens 100
-    }
-    model_id "anthropic.claude-3-haiku-20240307-v1:0"
-    // model_id "meta.llama3-8b-instruct-v1:0"
-    // model_id "mistral.mistral-7b-instruct-v0:2"
-    api_key ""
+    model claude-3-haiku-20240307
+    api_key env.ANTHROPIC_API_KEY
+    max_tokens 1000
   }
 }
 
+// client<llm> AwsBedrock {
+//   provider aws-bedrock
+//   options {
+//     inference_configuration {
+//       max_tokens 100
+//     }
+//     model_id "anthropic.claude-3-haiku-20240307-v1:0"
+//     // model_id "meta.llama3-8b-instruct-v1:0"
+//     // model_id "mistral.mistral-7b-instruct-v0:2"
+//     api_key ""
+//   }
+// }
+
 client<llm> Claude {
   provider anthropic
   options {