BoundaryML
diff --git a/‎.github/workflows/primary.yml‎
Lines changed: 2 additions & 0 deletions b/‎.github/workflows/primary.yml‎
Lines changed: 2 additions & 0 deletions
diff --git a/‎engine/Cargo.lock‎
Lines changed: 1 addition & 5 deletions b/‎engine/Cargo.lock‎
Lines changed: 1 addition & 5 deletions
diff --git a/‎engine/baml-lib/baml-types/src/baml_value.rs‎
Lines changed: 8 additions & 0 deletions b/‎engine/baml-lib/baml-types/src/baml_value.rs‎
Lines changed: 8 additions & 0 deletions
diff --git a/‎engine/baml-lib/baml-types/src/tracing/events.rs‎
Lines changed: 10 additions & 6 deletions b/‎engine/baml-lib/baml-types/src/tracing/events.rs‎
Lines changed: 10 additions & 6 deletions
diff --git a/‎engine/baml-runtime/src/tracingv2/storage/storage.rs‎
Lines changed: 38 additions & 13 deletions b/‎engine/baml-runtime/src/tracingv2/storage/storage.rs‎
Lines changed: 38 additions & 13 deletions
@@ -218,6 +218,8 @@ jobs:
           # Find Ruby library path and add it to LD_LIBRARY_PATH
           RUBY_LIB_PATH=$(ruby -e 'puts RbConfig::CONFIG["libdir"]')
           export LD_LIBRARY_PATH="${RUBY_LIB_PATH}:${LD_LIBRARY_PATH}"
+          # We need to build baml_cffi first so that the baml_cffi.dylib is available for the tests
+          cargo build --package baml_cffi
           cargo test --lib
         working-directory: engine
 
 
@@ -30,6 +30,14 @@ pub enum BamlValue {
     Null,
 }
 
+impl TryFrom<serde_json::Value> for BamlValue {
+    type Error = anyhow::Error;
+
+    fn try_from(value: serde_json::Value) -> Result<Self, Self::Error> {
+        Ok(serde_json::from_value(value)?)
+    }
+}
+
 impl serde::Serialize for BamlValue {
     fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
         match self {
 
@@ -379,7 +379,7 @@ where
     Option::<HashMap<String, String>>::deserialize(deserializer)
 }
 
-#[derive(Debug, Serialize, Deserialize)]
+#[derive(Debug, Serialize, Deserialize, Clone)]
 pub struct HTTPRequest {
     // since LLM requests could be made in parallel, we need to match the response to the request
     pub id: HttpRequestId,
@@ -439,17 +439,17 @@ pub struct ClientDetails {
     pub options: IndexMap<String, serde_json::Value>,
 }
 
-#[derive(Debug, Serialize, Deserialize)]
+#[derive(Clone, Debug, Serialize, Deserialize)]
 pub struct HTTPResponse {
     // since LLM requests could be made in parallel, we need to match the response to the request
     pub request_id: HttpRequestId,
     pub status: u16,
     #[serde(serialize_with = "serialize_redacted_optional_headers")]
     #[serde(deserialize_with = "deserialize_optional_headers")]
     headers: Option<HashMap<String, String>>,
-    pub body: HTTPBody,
+    pub body: Arc<HTTPBody>,
 
-    pub client_details: ClientDetails,
+    pub client_details: Arc<ClientDetails>,
 }
 
 #[derive(Debug, Serialize, Deserialize)]
@@ -470,11 +470,15 @@ impl HTTPResponse {
             request_id,
             status,
             headers,
-            body,
-            client_details,
+            body: Arc::new(body),
+            client_details: Arc::new(client_details),
         }
     }
 
+    pub fn id(&self) -> &HttpRequestId {
+        &self.request_id
+    }
+
     pub fn headers(&self) -> Option<&HashMap<String, String>> {
         self.headers.as_ref()
     }
 
@@ -327,16 +327,22 @@ fn build_function_log(
 
             // Streaming call
             calls.push(LLMCallKind::Stream(LLMStreamCall {
-                client_name: client,
-                provider,
+                llm_call: LLMCall {
+                    client_name: client,
+                    provider,
+                    timing: Timing {
+                        start_time_utc_ms: start_t,
+                        duration_ms: Some(partial_duration),
+                    },
+                    request: call_acc.http_request.clone(),
+                    response: call_acc.http_response.clone(),
+                    usage: Some(local_usage),
+                    selected: call_acc.llm_response.is_some(),
+                },
                 timing: StreamTiming {
                     start_time_utc_ms: start_t,
                     duration_ms: Some(partial_duration),
                 },
-                request: call_acc.http_request.clone(),
-                response: call_acc.http_response.clone(),
-                usage: Some(local_usage),
-                selected: call_acc.llm_response.is_some(),
                 sse_chunks,
             }));
         }
@@ -568,7 +574,21 @@ impl LLMCallKind {
     pub fn selected(&self) -> bool {
         match self {
             LLMCallKind::Basic(c) => c.selected,
-            LLMCallKind::Stream(c) => c.selected,
+            LLMCallKind::Stream(c) => c.llm_call.selected,
+        }
+    }
+
+    pub fn as_request(&self) -> Option<&LLMCall> {
+        match self {
+            LLMCallKind::Basic(c) => Some(c),
+            LLMCallKind::Stream(c) => None,
+        }
+    }
+
+    pub fn as_stream(&self) -> Option<&LLMStreamCall> {
+        match self {
+            LLMCallKind::Basic(c) => None,
+            LLMCallKind::Stream(c) => Some(c),
         }
     }
 }
@@ -586,13 +606,8 @@ pub struct LLMCall {
 
 #[derive(Debug, Default, Clone, Serialize)]
 pub struct LLMStreamCall {
-    pub client_name: String,
-    pub provider: String,
+    pub llm_call: LLMCall,
     pub timing: StreamTiming,
-    pub request: Option<Arc<HTTPRequest>>,
-    pub response: Option<Arc<HTTPResponse>>,
-    pub usage: Option<Usage>,
-    pub selected: bool,
     pub sse_chunks: Option<Arc<LLMHTTPStreamResponse>>,
 }
 
@@ -640,6 +655,16 @@ impl Collector {
         }
     }
 
+    pub fn clear(&self) -> usize {
+        let mut guard = self.tracked_ids.lock().unwrap();
+        for fid in guard.iter() {
+            BAML_TRACER.lock().unwrap().dec_ref(fid);
+        }
+        let len = guard.len();
+        guard.clear();
+        len
+    }
+
     pub fn function_logs(&self) -> Vec<FunctionLog> {
         let guard = self.tracked_ids.lock().unwrap();
         guard