vscode: add link to tracing, show token counts (#703)

* vscode + fiddle: add Boundary Studio link to traced tests * vscode + fiddle: show input and output token counts * runtime: surface token stats for openai streams, now that they support it
BoundaryML · Jun 22, 2024 · 64aa18a · 64aa18a
1 parent 21b1209
commit 64aa18a
Show file tree

Hide file tree

Showing 19 changed files with 8,536 additions and 6,980 deletions.
diff --git a/engine/.vscode/settings.json b/engine/.vscode/settings.json
diff --git a/engine/Cargo.lock b/engine/Cargo.lock
diff --git a/engine/Cargo.toml b/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" }

diff --git a/engine/baml-runtime/src/internal/llm_client/primitive/openai/openai_client.rs b/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())))
                     },

diff --git a/engine/baml-schema-wasm/.gitignore b/engine/baml-schema-wasm/.gitignore
@@ -1,2 +1,3 @@
+dist/src/*
 nodejs/src/*
 web/src/*
diff --git a/engine/baml-schema-wasm/Cargo.toml b/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"

diff --git a/engine/baml-schema-wasm/dist/src/baml_schema_build.d.ts b/engine/baml-schema-wasm/dist/src/baml_schema_build.d.ts
diff --git a/engine/baml-schema-wasm/dist/src/baml_schema_build.js b/engine/baml-schema-wasm/dist/src/baml_schema_build.js
diff --git a/engine/baml-schema-wasm/dist/src/baml_schema_build_bg.wasm b/engine/baml-schema-wasm/dist/src/baml_schema_build_bg.wasm
diff --git a/engine/baml-schema-wasm/dist/src/baml_schema_build_bg.wasm.d.ts b/engine/baml-schema-wasm/dist/src/baml_schema_build_bg.wasm.d.ts