vscode: add link to tracing, show token counts

engine/Cargo.toml

@@ -48,6 +48,7 @@ serde = { version = "1", features = ["derive"] }
 static_assertions = "1.1.0"
 strum = { version = "0.26.2", features = ["derive"] }
 strum_macros = "0.26.2"
+time = { version = "0.3.36", features = ["formatting"] }
 walkdir = "2.5.0"
 web-time = "1.1.0"
 baml-types = { path = "baml-lib/baml-types" }

engine/baml-runtime/src/internal/llm_client/primitive/openai/openai_client.rs

@@ -272,7 +272,13 @@ impl RequestBuilder for OpenAIClient {
         }
 
         if stream {
-            body_obj.insert("stream".into(), true.into());
+            body_obj.insert("stream".into(), json!(true));
+            body_obj.insert(
+                "stream_options".into(),
+                json!({
+                    "include_usage": true,
+                }),
+            );
         }
 
         req.json(&body)
@@ -366,6 +372,11 @@ impl SseResponseTrait for OpenAIClient {
                             }
                         }
                         inner.latency = instant_start.elapsed();
+                        if let Some(usage) = event.usage.as_ref() {
+                            inner.metadata.prompt_tokens = Some(usage.prompt_tokens);
+                            inner.metadata.output_tokens = Some(usage.completion_tokens);
+                            inner.metadata.total_tokens = Some(usage.total_tokens);
+                        }
 
                         std::future::ready(Some(LLMResponse::Success(inner.clone())))
                     },

engine/baml-schema-wasm/.gitignore

@@ -1,2 +1,3 @@
+dist/src/*
 nodejs/src/*
 web/src/*

engine/baml-schema-wasm/Cargo.toml

@@ -25,6 +25,7 @@ log.workspace = true
 serde.workspace = true
 serde_json.workspace = true
 serde-wasm-bindgen = "0.4"
+time.workspace = true
 uuid = { version = "1.8", features = ["v4", "js"] }
 wasm-bindgen = "=0.2.92"
 wasm-bindgen-futures = "0.4.42"