Add ruby collector (#1587)

<!-- ELLIPSIS_HIDDEN -->


> [!IMPORTANT]
> Add Ruby collector to BAML for enhanced tracing and logging, updating
FFI, client templates, and tests.
> 
>   - **Behavior**:
> - Adds `Collector` class to `log_collector.rs` for managing function
logs and usage.
> - Updates `BamlRuntimeFfi` in `lib.rs` to handle collectors in
`call_function` and `stream_function`.
> - Modifies `client.rb.j2` to support `collector` in `baml_options`.
>   - **Models**:
> - Introduces `FunctionLog`, `Usage`, `Timing`, `StreamTiming`,
`LLMCall`, and `LLMStreamCall` in `log_collector.rs`.
> - Implements `lang_wrapper!` macro in `lang_wrapper.rs` for wrapping
types.
>   - **Dependencies**:
>     - Updates `Cargo.toml` to use `serde_magnus` from the registry.
>   - **Tests**:
>     - Adds `test_collector.rb` for testing collector functionality.
> - Updates `README.md` in `integ-tests/ruby` with test instructions.
>   - **Misc**:
>     - Registers new classes in `init` function in `lib.rs`.
> - Adds `define_all_in_ruby` in `log_collector.rs` to register Ruby
classes.
> 
> <sup>This description was created by </sup>[<img alt="Ellipsis"
src="https://img.shields.io/badge/Ellipsis-blue?color=175173">](https://www.ellipsis.dev?ref=BoundaryML%2Fbaml&utm_source=github&utm_medium=referral)<sup>
for dbfa014. It will automatically
update as commits are pushed.</sup>


<!-- ELLIPSIS_HIDDEN -->

Author: aaronvg

engine/Cargo.lock

@@ -5,19 +5,19 @@
 //! for a single FunctionId, even if multiple Collectors or FunctionLogs want it.
 //! It uses manual reference counting (`inc_ref` / `dec_ref`) to free memory for
 //! a FunctionId as soon as there are no more "owners."
+use baml_types::tracing::events::{
+    FunctionEnd, FunctionId, FunctionStart, HTTPRequest, HTTPResponse, LoggedLLMRequest,
+    LoggedLLMResponse, TraceData, TraceEvent,
+};
 use indexmap::{IndexMap, IndexSet};
 use once_cell::sync::Lazy;
+use serde::Serialize;
 use std::collections::{HashMap, HashSet};
 use std::fmt;
 use std::hash::Hash;
 use std::sync::{Arc, Mutex};
 use uuid::Uuid;
 
-use baml_types::tracing::events::{
-    FunctionEnd, FunctionId, FunctionStart, HTTPRequest, HTTPResponse, LoggedLLMRequest,
-    LoggedLLMResponse, TraceData, TraceEvent,
-};
-
 use crate::tracingv2::publisher::publisher::PublisherMessage;
 
 use super::super::publisher::PUBLISHING_CHANNEL;
@@ -470,7 +470,7 @@ impl Drop for FunctionLog {
 /// Represents the "inner" data for a single function call
 /// (the real set of usage/calls/timing, etc.).
 ///
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, Serialize)]
 pub struct FunctionLogInner {
     pub id: String,
     pub function_name: String,
@@ -489,28 +489,29 @@ impl FunctionLogInner {
     }
 }
 
-#[derive(Debug, Default, Clone, Hash, Eq, PartialEq)]
+#[derive(Debug, Default, Clone, Hash, Eq, PartialEq, Serialize)]
 pub struct Usage {
     pub input_tokens: Option<i64>,
     pub output_tokens: Option<i64>,
 }
 
-#[derive(Debug, Default, Clone, Hash, Eq, PartialEq)]
+#[derive(Debug, Default, Clone, Hash, Eq, PartialEq, Serialize)]
 pub struct Timing {
     pub start_time_utc_ms: i64,
     pub duration_ms: Option<i64>,
     pub time_to_first_parsed_ms: Option<i64>,
 }
 
-#[derive(Debug, Default, Clone, Hash, Eq, PartialEq)]
+#[derive(Debug, Default, Clone, Hash, Eq, PartialEq, Serialize)]
 pub struct StreamTiming {
     pub start_time_utc_ms: i64,
     pub duration_ms: Option<i64>,
     pub time_to_first_parsed_ms: Option<i64>,
     pub time_to_first_token_ms: Option<i64>,
 }
 
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, Serialize)]
+#[serde(untagged)]
 pub enum LLMCallKind {
     Basic(LLMCall),
     Stream(LLMStreamCall),
@@ -526,7 +527,7 @@ impl LLMCallKind {
     }
 }
 
-#[derive(Debug, Default, Clone)]
+#[derive(Debug, Default, Clone, Serialize)]
 pub struct LLMCall {
     pub client_name: String,
     pub provider: String,
@@ -537,7 +538,7 @@ pub struct LLMCall {
     pub selected: bool,
 }
 
-#[derive(Debug, Default, Clone)]
+#[derive(Debug, Default, Clone, Serialize)]
 pub struct LLMStreamCall {
     pub client_name: String,
     pub provider: String,

engine/baml-runtime/src/tracingv2/storage/storage.rs

@@ -51,7 +51,7 @@ module Baml
         {% for (name, type) in fn.args -%}
         {{name}}: {{type}},
         {%- endfor %}
-        baml_options: T::Hash[Symbol, T.any(Baml::TypeBuilder, Baml::ClientRegistry)]
+        baml_options: T::Hash[Symbol, T.any(Baml::TypeBuilder, Baml::ClientRegistry, T.any(Baml::Collector, T::Array[Baml::Collector]))]
       ).returns({{ fn.return_type }})
     }
     def {{fn.name}}(
@@ -65,8 +65,14 @@ module Baml
         {# We rely on sorbet-runtime to give errors about the list of allowed kwargs #}
         raise ArgumentError.new("{{fn.name}} may only be called with keyword arguments")
       end
-      if (baml_options.keys - [:client_registry, :tb]).any?
-        raise ArgumentError.new("Received unknown keys in baml_options (valid keys: :client_registry, :tb): #{baml_options.keys - [:client_registry, :tb]}")
+      if (baml_options.keys - [:client_registry, :tb, :collector]).any?
+        raise ArgumentError.new("Received unknown keys in baml_options (valid keys: :client_registry, :tb, :collector): #{baml_options.keys - [:client_registry, :tb, :collector]}")
+      end
+
+      collector = if baml_options[:collector]
+        baml_options[:collector].is_a?(Array) ? baml_options[:collector] : [baml_options[:collector]]
+      else
+        []
       end
 
       raw = @runtime.call_function(
@@ -79,6 +85,7 @@ module Baml
         @ctx_manager,
         baml_options[:tb]&.instance_variable_get(:@registry),
         baml_options[:client_registry],
+        collector,
       )
       (raw.parsed_using_types(Baml::Types, Baml::PartialTypes, false))
     end
@@ -103,7 +110,7 @@ module Baml
         {% for (name, type) in fn.args -%}
         {{name}}: {{type}},
         {%- endfor %}
-        baml_options: T::Hash[Symbol, T.any(Baml::TypeBuilder, Baml::ClientRegistry)]
+        baml_options: T::Hash[Symbol, T.any(Baml::TypeBuilder, Baml::ClientRegistry, T.any(Baml::Collector, T::Array[Baml::Collector]))]
       ).returns(Baml::BamlStream[{{ fn.return_type }}])
     }
     def {{fn.name}}(
@@ -117,8 +124,14 @@ module Baml
         {# We rely on sorbet-runtime to give errors about the list of allowed kwargs #}
         raise ArgumentError.new("{{fn.name}} may only be called with keyword arguments")
       end
-      if (baml_options.keys - [:client_registry, :tb]).any?
-        raise ArgumentError.new("Received unknown keys in baml_options (valid keys: :client_registry, :tb): #{baml_options.keys - [:client_registry, :tb]}")
+      if (baml_options.keys - [:client_registry, :tb, :collector]).any?
+        raise ArgumentError.new("Received unknown keys in baml_options (valid keys: :client_registry, :tb, :collector): #{baml_options.keys - [:client_registry, :tb, :collector]}")
+      end
+
+      collector = if baml_options[:collector]
+        baml_options[:collector].is_a?(Array) ? baml_options[:collector] : [baml_options[:collector]]
+      else
+        []
       end
 
       raw = @runtime.stream_function(
@@ -131,6 +144,7 @@ module Baml
         @ctx_manager,
         baml_options[:tb]&.instance_variable_get(:@registry),
         baml_options[:client_registry],
+        collector,
       )
       Baml::BamlStream[{{fn.partial_return_type}}, {{fn.return_type}}].new(
         ffi_stream: raw,

engine/language_client_codegen/src/ruby/templates/client.rb.j2

@@ -1,7 +1,7 @@
 PATH
   remote: .
   specs:
-    baml (0.76.2)
+    baml (0.78.0)
 
 GEM
   remote: https://rubygems.org/

engine/language_client_ruby/Gemfile.lock

@@ -30,7 +30,7 @@ rb-sys = { version = "0.9.99", features = [
 ] }
 serde.workspace = true
 serde_json.workspace = true
-serde_magnus = { git = "https://github.com/BoundaryML/serde-magnus.git", branch = "sam/magnus-0.7.1" }
+serde_magnus = { version = "0.9.0" }
 tokio = { version = "1", features = ["full"] }
 tracing-subscriber = { version = "0.3.18", features = ["json", "env-filter","valuable"] }
 

engine/language_client_ruby/ext/ruby_ffi/Cargo.toml

@@ -1,10 +1,11 @@
 use baml_runtime::BamlRuntime;
 use baml_types::BamlValue;
-use magnus::{class, function, method, prelude::*, Error, RHash, Ruby};
+use magnus::{class, function, method, prelude::*, Error, RArray, RHash, Ruby};
 use std::collections::HashMap;
 use std::path::PathBuf;
 use std::sync::Arc;
 use tracing_subscriber::EnvFilter;
+use types::log_collector::Collector;
 
 use function_result::FunctionResult;
 use function_result_stream::FunctionResultStream;
@@ -112,6 +113,7 @@ impl BamlRuntimeFfi {
         ctx: &RuntimeContextManager,
         type_registry: Option<&types::type_builder::TypeBuilder>,
         client_registry: Option<&types::client_registry::ClientRegistry>,
+        collector: RArray,
     ) -> Result<FunctionResult> {
         let args = match ruby_to_json::RubyToJson::convert_hash_to_json(args) {
             Ok(args) => args.into_iter().collect(),
@@ -123,14 +125,18 @@ impl BamlRuntimeFfi {
             }
         };
 
+        let mut collectors = Vec::new();
+        for i in collector.into_iter() {
+            collectors.push(<&Collector>::try_convert(i)?.inner.clone());
+        }
+
         let retval = match rb_self.t.block_on(rb_self.inner.call_function(
             function_name.clone(),
             &args,
             &ctx.inner,
             type_registry.map(|t| &t.inner),
             client_registry.map(|c| c.inner.borrow_mut()).as_deref(),
-            // TODO: wire this
-            Some(vec![]),
+            Some(collectors),
         )) {
             (Ok(res), _) => Ok(FunctionResult::new(res)),
             (Err(e), _) => Err(Error::new(
@@ -153,6 +159,7 @@ impl BamlRuntimeFfi {
         ctx: &RuntimeContextManager,
         type_registry: Option<&types::type_builder::TypeBuilder>,
         client_registry: Option<&types::client_registry::ClientRegistry>,
+        collector: RArray,
     ) -> Result<FunctionResultStream> {
         let args = match ruby_to_json::RubyToJson::convert_hash_to_json(args) {
             Ok(args) => args.into_iter().collect(),
@@ -164,6 +171,11 @@ impl BamlRuntimeFfi {
             }
         };
 
+        let mut collectors = Vec::new();
+        for i in collector.into_iter() {
+            collectors.push(<&Collector>::try_convert(i)?.inner.clone());
+        }
+
         log::debug!("Streaming {function_name} with:\nargs: {args:#?}\nctx ???");
 
         let retval = match rb_self.inner.stream_function(
@@ -172,8 +184,7 @@ impl BamlRuntimeFfi {
             &ctx.inner,
             type_registry.map(|t| &t.inner),
             client_registry.map(|c| c.inner.borrow_mut()).as_deref(),
-            // TODO: wire this
-            Some(vec![]),
+            Some(collectors),
         ) {
             Ok(res) => Ok(FunctionResultStream::new(res, rb_self.t.clone())),
             Err(e) => Err(Error::new(
@@ -255,10 +266,10 @@ fn init(ruby: &Ruby) -> Result<()> {
         "create_context_manager",
         method!(BamlRuntimeFfi::create_context_manager, 0),
     )?;
-    runtime_class.define_method("call_function", method!(BamlRuntimeFfi::call_function, 5))?;
+    runtime_class.define_method("call_function", method!(BamlRuntimeFfi::call_function, 6))?;
     runtime_class.define_method(
         "stream_function",
-        method!(BamlRuntimeFfi::stream_function, 5),
+        method!(BamlRuntimeFfi::stream_function, 6),
     )?;
 
     FunctionResult::define_in_ruby(&module)?;
@@ -277,6 +288,9 @@ fn init(ruby: &Ruby) -> Result<()> {
     types::media::Audio::define_in_ruby(&module)?;
     types::media::Image::define_in_ruby(&module)?;
 
+    // Register the new log collector classes
+    types::log_collector::define_all_in_ruby(&module)?;
+
     // everything below this is for our own testing purposes
     module.define_module_function(
         "roundtrip",

engine/language_client_ruby/ext/ruby_ffi/src/lib.rs

@@ -28,4 +28,33 @@ macro_rules! lang_wrapper {
         }
     };
 
+
+    ($name:ident, $wrap_name:expr, $type:ty, clone_safe $(, $attr_name:ident : $attr_type:ty)*) => {
+        #[magnus::wrap(class = $wrap_name, free_immediately, size)]
+        pub struct $name {
+            pub(crate) inner: std::sync::Arc<$type>,
+            $($attr_name: $attr_type),*
+        }
+
+        impl From<$type> for $name {
+            fn from(inner: $type) -> Self {
+                Self {
+                    inner: std::sync::Arc::new(inner),
+                    $($attr_name: Default::default()),*
+                }
+            }
+        }
+
+
+        impl From<std::sync::Arc<$type>> for $name {
+            fn from(inner: std::sync::Arc<$type>) -> Self {
+                Self {
+                    inner,
+                    $($attr_name: Default::default()),*
+                }
+            }
+        }
+    };
+
+
 }