hellovai/tracing

.mise.toml

@@ -2,4 +2,4 @@
 node = "20.14"
 ruby = "3.1"
 pnpm = "9.9"
-poetry = "1.8.4"
+poetry = "1.8.5"

engine/Cargo.toml

@@ -5,11 +5,13 @@ members = [
   "baml-runtime",
   "baml-schema-wasm",
   "bstd",
+  "btrace",
   "cli",
   "language_client_codegen",
   "language_client_python",
   "language_client_ruby/ext/ruby_ffi",
   "language_client_typescript",
+  "sandbox",
 ]
 default-members = [
   "baml-lib/*",
@@ -20,6 +22,7 @@ default-members = [
   # also have something to do with resolver=2 and how duplicate dep builds are handled)
   # "baml-schema-wasm",
   "bstd",
+  "btrace",
   "cli",
   "language_client_codegen",
   "language_client_python",
@@ -32,8 +35,10 @@ anyhow = "1.0"
 array-macro = "2"
 askama = "0.12.1"
 baml-cli = { path = "cli" }
+baml-types = { path = "baml-lib/baml-types" }
 base64 = "0.22.1"
 bstd = { path = "bstd" }
+btrace = { path = "btrace" }
 bytes = "1.6.0"
 cfg-if = "1.0.0"
 clap = { version = "4.4.6", features = ["cargo", "derive"] }
@@ -47,6 +52,12 @@ http = "1.1.0"
 http-body = "1.0.0"
 indexmap = { version = "2.1.0", features = ["serde"] }
 indoc = "2.0.5"
+jsonish = { path = "baml-lib/jsonish" }
+internal-baml-codegen = { path = "language_client_codegen" }
+internal-baml-core = { path = "baml-lib/baml-core" }
+internal-baml-jinja = { path = "baml-lib/jinja-runtime" }
+internal-baml-schema-ast = { path = "baml-lib/schema-ast" }
+internal-llm-client = { path = "baml-lib/llm-client" }
 log = "0.4.20"
 # TODO: disable imports, etc
 minijinja = { version = "1.0.16", default-features = false, features = [
@@ -67,7 +78,8 @@ minijinja = { version = "1.0.16", default-features = false, features = [
   # loader
   #
 ] }
-pretty_assertions = "1.4.1"
+once_cell = "1"
+pretty_assertions = "1.4.0"
 rand = "0.8.5"
 regex = "1.10.4"
 reqwest = { version = "0.12.5", features = [
@@ -83,16 +95,13 @@ static_assertions = "1.1.0"
 strum = { version = "0.26.2", features = ["derive"] }
 strum_macros = "0.26.2"
 time = { version = "0.3.36", features = ["formatting"] }
+tracing-core = { version = "0.1.31" }
 pin-project-lite = "0.2.14"
 url = "2.5.2"
+uuid = { version = "1.8.0", features = ["v4", "v7", "serde"] }
 walkdir = "2.5.0"
 web-time = "1.1.0"
 
-baml-types = { path = "baml-lib/baml-types" }
-internal-baml-codegen = { path = "language_client_codegen" }
-internal-baml-core = { path = "baml-lib/baml-core" }
-internal-baml-jinja = { path = "baml-lib/jinja" }
-internal-baml-schema-ast = { path = "baml-lib/schema-ast" }
 
 [workspace.package]
 version = "0.72.0"

engine/baml-lib/baml-types/Cargo.toml

@@ -11,16 +11,21 @@ dead_code = "deny"
 unused_imports = "deny"
 unused_variables = "deny"
 
+# Attention:
+#   Please try to keep this crate low-dependency! We want to keep this as minimal as possible.
 [dependencies]
 anyhow.workspace = true
 clap.workspace = true
 derive_builder.workspace = true
+log.workspace = true
+minijinja.workspace = true
+once_cell.workspace = true
 serde.workspace = true
 serde_json.workspace = true
 strum.workspace = true
-minijinja.workspace = true
-log.workspace = true
-once_cell = "1"
+tracing-core.workspace = true
+web-time.workspace = true
+tokio.workspace = true
 
 [dependencies.indexmap]
 workspace = true

engine/baml-lib/baml-types/src/lib.rs

@@ -2,6 +2,7 @@ mod constraint;
 mod map;
 mod media;
 mod minijinja;
+pub mod tracing;
 
 mod baml_value;
 mod field_type;

engine/baml-lib/baml-types/src/tracing/events.rs

@@ -0,0 +1,131 @@
+use anyhow::Result;
+use crate::BamlValue;
+
+// TODO: use a prefixed UUID type for this
+type SpanId = String;
+
+#[derive(Clone, PartialEq, Eq, Hash)]
+pub struct FunctionId(pub SpanId);
+
+#[derive(Clone, PartialEq, Eq, Hash)]
+pub struct ContentId(pub SpanId);
+
+// THESE ARE NOT CLONEABLE!!
+pub struct LogEvent {
+    /*
+     * (span_id, content_span_id) is a unique identifier for a log event
+     * The query (span_id, *) gets all logs for a function call
+     */
+
+    pub span_id: FunctionId,
+    pub content_span_id: ContentId,
+
+    // The chain of spans that lead to this log event
+    // Includes span_id at the last position (content_span_id is not included)
+    pub span_chain: Vec<FunctionId>,
+
+    // The timestamp of the log
+    pub timestamp: web_time::Instant,
+    // The content of the log
+    pub content: LogEventContent,
+}
+
+
+pub enum LogEventContent {
+    // All start events
+    FunctionStart(FunctionStart),
+    // All end events
+    FunctionEnd(FunctionEnd),
+
+    // The rest are intermediate events that happen between start and end
+
+    // LLM request
+    LLMRequest(LLMRequest),
+    // Raw HTTP request to the LLM
+    RawLLMRequest(HTTPRequest),
+
+    // Do to streaming, its possible to have multiple responses for a single request
+    // ----
+    // Raw HTTP response from the LLM
+    RawLLMResponse(HTTPResponse),
+    // Parsed LLM response
+    LLMResponse(Result<LLMResponse>),
+    // ----
+
+    // We don't want to store the parsed LLM response in the log event
+    // as we have it in FunctionEnd
+    Parsed(Result<()>),
+}
+
+pub struct BamlOptions {
+    pub type_builder: Option<serde_json::Value>,
+    pub client_registry: Option<serde_json::Value>,
+}
+
+pub struct FunctionStart {
+    pub name: String,
+    pub args: Vec<BamlValue>,
+    pub options: BamlOptions,
+}
+
+pub struct FunctionEnd {
+    pub result: Result<BamlValue>,
+    // Everything below is duplicated from the start event
+    // to deal with the case where the log is dropped.
+    // P2: as we can for now assume logs are not dropped,
+
+    // pub name: String,
+    // pub start_timestamp: web_time::Instant,
+    // pub start_args: Vec<BamlValue>,
+}
+
+// LLM specific events
+
+// TODO: fix this.
+pub type Prompt = serde_json::Value;
+
+pub enum LLMClient {
+    Ref(String),
+    ShortHand(String, String),
+}
+
+pub struct LLMRequest {
+    pub client: LLMClient,
+    pub params: serde_json::Value,
+    pub prompt: Prompt,
+}
+
+pub struct HTTPRequest {
+    // since LLM requests could be made in parallel, we need to match the response to the request
+    pub request_id: ContentId,
+    pub url: String,
+    pub method: String,
+    pub headers: serde_json::Value,
+    pub body: serde_json::Value,
+}
+
+pub struct HTTPResponse {
+    // since LLM requests could be made in parallel, we need to match the response to the request
+    pub request_id: ContentId,
+    pub status: u16,
+    pub headers: serde_json::Value,
+    pub body: serde_json::Value,
+}
+
+
+pub struct LLMResponse {
+    // since LLM requests could be made in parallel, we need to match the response to the request
+    pub request_id: ContentId,
+    // Fully qualified model name
+    pub finish_reason: String,
+    pub model: String,
+    pub usage: LLMUsage,
+    pub string_output: String,
+}
+
+pub struct LLMUsage {
+    pub input_tokens: Option<u64>,
+    pub output_tokens: Option<u64>,
+    pub total_tokens: Option<u64>,
+}
+