braintrust 0.0.5 → 0.0.6

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 6321acf7b780922ed97ea3cc57dde47a52947a10650a082dcfd9af780056d99a
-  data.tar.gz: 67c181e53537829931de704c7503cc056646652f9c1a61d914bc1ee0b7af69a2
+  metadata.gz: 866cb2e797502f00cda1625ad90f4d734b4b83f0d21d8243675a933fae9df693
+  data.tar.gz: f74151b0e18b12cf19b61b1b75b2f58e784d4171f21c0996526d29c719174260
 SHA512:
-  metadata.gz: bb8546fdbf0a448016a1d31ceb8729a40be59e0d8d081ef275f763a11dbb2f5df0134ec52fc3b1c15c41d9dcdf42fbbe6becaf00ab7ac882c8f2f7e173a9a61f
-  data.tar.gz: 41e6d13504302a3b3ec26697cb50ce4736040d910c8e293d37088311922daa77274fc4214c86543ca21209bf56e82b2f6f00d5fbc2d7d0d0baad6f1e77cc48ff
+  metadata.gz: ad2f68a6de8d547b6a609c3393522c4ae3dfcb441a9fc841484bbbcb21de7648da7a00cd625612d98c6b99e4ad41186a2bc3fff706e17b9797e7ac514e685923
+  data.tar.gz: f0613e5fa08c07333c74467ec7830a40f72905475e35becf7a2add077168c7554046aa9a3824fe24006870338163526e8d170cfd25727af5d53416283ae03714

data/README.md

@@ -155,7 +155,7 @@ message = tracer.in_span("chat-message") do |span|
   root_span = span
 
   client.messages.create(
-    model: "claude-3-5-sonnet-20241022",
+    model: "claude-3-haiku-20240307",
     max_tokens: 100,
     system: "You are a helpful assistant.",
     messages: [
@@ -171,6 +171,34 @@ puts "View trace at: #{Braintrust::Trace.permalink(root_span)}"
 OpenTelemetry.tracer_provider.shutdown
 ```
 
+### RubyLLM Tracing
+
+```ruby
+require "braintrust"
+require "ruby_llm"
+
+Braintrust.init
+
+# Wrap RubyLLM globally (wraps all Chat instances)
+Braintrust::Trace::Contrib::Github::Crmne::RubyLLM.wrap
+
+tracer = OpenTelemetry.tracer_provider.tracer("ruby-llm-app")
+root_span = nil
+
+response = tracer.in_span("chat") do |span|
+  root_span = span
+
+  chat = RubyLLM.chat(model: "gpt-4o-mini")
+  chat.ask("Say hello!")
+end
+
+puts "Response: #{response.content}"
+
+puts "View trace at: #{Braintrust::Trace.permalink(root_span)}"
+
+OpenTelemetry.tracer_provider.shutdown
+```
+
 ### Attachments
 
 Attachments allow you to log binary data (images, PDFs, audio, etc.) as part of your traces. This is particularly useful for multimodal AI applications like vision models.
@@ -236,7 +264,9 @@ Check out the [`examples/`](./examples/) directory for complete working examples
 - [eval.rb](./examples/eval.rb) - Create and run evaluations with custom test cases and scoring functions
 - [trace.rb](./examples/trace.rb) - Manual span creation and tracing
 - [openai.rb](./examples/openai.rb) - Automatically trace OpenAI API calls
+- [alexrudall_openai.rb](./examples/alexrudall_openai.rb) - Automatically trace ruby-openai gem API calls
 - [anthropic.rb](./examples/anthropic.rb) - Automatically trace Anthropic API calls
+- [ruby_llm.rb](./examples/ruby_llm.rb) - Automatically trace RubyLLM API calls
 - [trace/trace_attachments.rb](./examples/trace/trace_attachments.rb) - Log attachments (images, PDFs) in traces
 - [eval/dataset.rb](./examples/eval/dataset.rb) - Run evaluations using datasets stored in Braintrust
 - [eval/remote_functions.rb](./examples/eval/remote_functions.rb) - Use remote scoring functions

data/lib/braintrust/state.rb

@@ -49,6 +49,20 @@ module Braintrust
       )
     end
 
+    # Create a State object directly with explicit parameters
+    # @param api_key [String] Braintrust API key (required)
+    # @param org_name [String, nil] Organization name
+    # @param org_id [String, nil] Organization ID (if provided, skips login - useful for testing)
+    # @param default_project [String, nil] Default project name
+    # @param app_url [String, nil] App URL (default: https://www.braintrust.dev)
+    # @param api_url [String, nil] API URL
+    # @param proxy_url [String, nil] Proxy URL
+    # @param blocking_login [Boolean] Login synchronously (default: false)
+    # @param enable_tracing [Boolean] Enable OpenTelemetry tracing (default: true)
+    # @param tracer_provider [TracerProvider, nil] Optional tracer provider
+    # @param config [Config, nil] Optional config object
+    # @param exporter [Exporter, nil] Optional exporter for testing
+    # @return [State] the created state
     def initialize(api_key: nil, org_name: nil, org_id: nil, default_project: nil, app_url: nil, api_url: nil, proxy_url: nil, blocking_login: false, enable_tracing: true, tracer_provider: nil, config: nil, exporter: nil)
       # Instance-level mutex for thread-safe login
       @login_mutex = Mutex.new
@@ -61,13 +75,17 @@ module Braintrust
       @app_url = app_url || "https://www.braintrust.dev"
       @api_url = api_url
       @proxy_url = proxy_url
-      @logged_in = false
       @config = config
 
-      # Perform login after state setup
-      if blocking_login
+      # If org_id is provided, we're already "logged in" (useful for testing)
+      # Otherwise, perform login to discover org info
+      if org_id
+        @logged_in = true
+      elsif blocking_login
+        @logged_in = false
         login
       else
+        @logged_in = false
         login_in_thread
       end
 

data/lib/braintrust/trace/contrib/anthropic.rb

@@ -2,6 +2,7 @@
 
 require "opentelemetry/sdk"
 require "json"
+require_relative "../tokens"
 
 module Braintrust
   module Trace
@@ -17,61 +18,11 @@ module Braintrust
         span.set_attribute(attr_name, JSON.generate(obj))
       end
 
-      # Parse usage tokens from Anthropic API response, handling cache tokens
-      # Maps Anthropic field names to Braintrust standard names:
-      # - input_tokens → contributes to prompt_tokens
-      # - cache_creation_input_tokens → prompt_cache_creation_tokens (and adds to prompt_tokens)
-      # - cache_read_input_tokens → prompt_cached_tokens (and adds to prompt_tokens)
-      # - output_tokens → completion_tokens
-      # - total_tokens → tokens (or calculated if missing)
-      #
+      # Parse usage tokens from Anthropic API response
       # @param usage [Hash, Object] usage object from Anthropic response
       # @return [Hash<String, Integer>] metrics hash with normalized names
       def self.parse_usage_tokens(usage)
-        metrics = {}
-        return metrics unless usage
-
-        # Convert to hash if it's an object
-        usage_hash = usage.respond_to?(:to_h) ? usage.to_h : usage
-
-        # Extract base values for calculation
-        input_tokens = 0
-        cache_creation_tokens = 0
-        cache_read_tokens = 0
-
-        usage_hash.each do |key, value|
-          next unless value.is_a?(Numeric)
-          key_str = key.to_s
-
-          case key_str
-          when "input_tokens"
-            input_tokens = value.to_i
-          when "cache_creation_input_tokens"
-            cache_creation_tokens = value.to_i
-            metrics["prompt_cache_creation_tokens"] = value.to_i
-          when "cache_read_input_tokens"
-            cache_read_tokens = value.to_i
-            metrics["prompt_cached_tokens"] = value.to_i
-          when "output_tokens"
-            metrics["completion_tokens"] = value.to_i
-          when "total_tokens"
-            metrics["tokens"] = value.to_i
-          else
-            # Keep other numeric fields as-is (future-proofing)
-            metrics[key_str] = value.to_i
-          end
-        end
-
-        # Calculate total prompt tokens (input + cache creation + cache read)
-        total_prompt_tokens = input_tokens + cache_creation_tokens + cache_read_tokens
-        metrics["prompt_tokens"] = total_prompt_tokens
-
-        # Calculate total tokens if not provided by Anthropic
-        if !metrics.key?("tokens") && metrics.key?("completion_tokens")
-          metrics["tokens"] = total_prompt_tokens + metrics["completion_tokens"]
-        end
-
-        metrics
+        Braintrust::Trace.parse_anthropic_usage_tokens(usage)
       end
 
       # Wrap an Anthropic::Client to automatically create spans for messages and responses

data/lib/braintrust/trace/contrib/github.com/alexrudall/ruby-openai/ruby-openai.rb

@@ -2,140 +2,134 @@
 
 require "opentelemetry/sdk"
 require "json"
+require_relative "../../../../tokens"
 
 module Braintrust
   module Trace
-    module AlexRudall
-      module RubyOpenAI
-        # Helper to safely set a JSON attribute on a span
-        # Only sets the attribute if obj is present
-        # @param span [OpenTelemetry::Trace::Span] the span to set attribute on
-        # @param attr_name [String] the attribute name (e.g., "braintrust.output_json")
-        # @param obj [Object] the object to serialize to JSON
-        # @return [void]
-        def self.set_json_attr(span, attr_name, obj)
-          return unless obj
-          span.set_attribute(attr_name, JSON.generate(obj))
-        end
+    module Contrib
+      module Github
+        module Alexrudall
+          module RubyOpenAI
+            # Helper to safely set a JSON attribute on a span
+            # Only sets the attribute if obj is present
+            # @param span [OpenTelemetry::Trace::Span] the span to set attribute on
+            # @param attr_name [String] the attribute name (e.g., "braintrust.output_json")
+            # @param obj [Object] the object to serialize to JSON
+            # @return [void]
+            def self.set_json_attr(span, attr_name, obj)
+              return unless obj
+              span.set_attribute(attr_name, JSON.generate(obj))
+            end
 
-        # Parse usage tokens from OpenAI API response, handling nested token_details
-        # Maps OpenAI field names to Braintrust standard names:
-        # - prompt_tokens → prompt_tokens
-        # - completion_tokens → completion_tokens
-        # - total_tokens → tokens
-        #
-        # @param usage [Hash] usage hash from OpenAI response
-        # @return [Hash<String, Integer>] metrics hash with normalized names
-        def self.parse_usage_tokens(usage)
-          metrics = {}
-          return metrics unless usage
-
-          # Basic token counts
-          metrics["prompt_tokens"] = usage["prompt_tokens"].to_i if usage["prompt_tokens"]
-          metrics["completion_tokens"] = usage["completion_tokens"].to_i if usage["completion_tokens"]
-          metrics["total_tokens"] = usage["total_tokens"].to_i if usage["total_tokens"]
-
-          # Rename total_tokens to tokens for consistency
-          metrics["tokens"] = metrics.delete("total_tokens") if metrics["total_tokens"]
-
-          metrics
-        end
+            # Parse usage tokens from OpenAI API response
+            # @param usage [Hash] usage hash from OpenAI response
+            # @return [Hash<String, Integer>] metrics hash with normalized names
+            def self.parse_usage_tokens(usage)
+              Braintrust::Trace.parse_openai_usage_tokens(usage)
+            end
 
-        # Wrap an OpenAI::Client (ruby-openai gem) to automatically create spans
-        # Supports both synchronous and streaming requests
-        # @param client [OpenAI::Client] the OpenAI client to wrap
-        # @param tracer_provider [OpenTelemetry::SDK::Trace::TracerProvider] the tracer provider (defaults to global)
-        def self.wrap(client, tracer_provider: nil)
-          tracer_provider ||= ::OpenTelemetry.tracer_provider
+            # Wrap an OpenAI::Client (ruby-openai gem) to automatically create spans
+            # Supports both synchronous and streaming requests
+            # @param client [OpenAI::Client] the OpenAI client to wrap
+            # @param tracer_provider [OpenTelemetry::SDK::Trace::TracerProvider] the tracer provider (defaults to global)
+            def self.wrap(client, tracer_provider: nil)
+              tracer_provider ||= ::OpenTelemetry.tracer_provider
 
-          # Wrap chat completions
-          wrap_chat(client, tracer_provider)
+              # Wrap chat completions
+              wrap_chat(client, tracer_provider)
 
-          client
-        end
+              client
+            end
 
-        # Wrap chat API
-        # @param client [OpenAI::Client] the OpenAI client
-        # @param tracer_provider [OpenTelemetry::SDK::Trace::TracerProvider] the tracer provider
-        def self.wrap_chat(client, tracer_provider)
-          # Create a wrapper module that intercepts the chat method
-          wrapper = Module.new do
-            define_method(:chat) do |parameters:|
-              tracer = tracer_provider.tracer("braintrust")
-
-              tracer.in_span("openai.chat.completions.create") do |span|
-                # Initialize metadata hash
-                metadata = {
-                  "provider" => "openai",
-                  "endpoint" => "/v1/chat/completions"
-                }
-
-                # Capture request metadata fields
-                metadata_fields = %w[
-                  model frequency_penalty logit_bias logprobs max_tokens n
-                  presence_penalty response_format seed service_tier stop
-                  stream stream_options temperature top_p top_logprobs
-                  tools tool_choice parallel_tool_calls user functions function_call
-                ]
-
-                metadata_fields.each do |field|
-                  field_sym = field.to_sym
-                  if parameters.key?(field_sym)
-                    # Special handling for stream parameter (it's a Proc)
-                    metadata[field] = if field == "stream"
-                      true  # Just mark as streaming
-                    else
-                      parameters[field_sym]
+            # Wrap chat API
+            # @param client [OpenAI::Client] the OpenAI client
+            # @param tracer_provider [OpenTelemetry::SDK::Trace::TracerProvider] the tracer provider
+            def self.wrap_chat(client, tracer_provider)
+              # Create a wrapper module that intercepts the chat method
+              wrapper = Module.new do
+                define_method(:chat) do |parameters:|
+                  tracer = tracer_provider.tracer("braintrust")
+
+                  tracer.in_span("openai.chat.completions.create") do |span|
+                    # Initialize metadata hash
+                    metadata = {
+                      "provider" => "openai",
+                      "endpoint" => "/v1/chat/completions"
+                    }
+
+                    # Capture request metadata fields
+                    metadata_fields = %w[
+                      model frequency_penalty logit_bias logprobs max_tokens n
+                      presence_penalty response_format seed service_tier stop
+                      stream stream_options temperature top_p top_logprobs
+                      tools tool_choice parallel_tool_calls user functions function_call
+                    ]
+
+                    metadata_fields.each do |field|
+                      field_sym = field.to_sym
+                      if parameters.key?(field_sym)
+                        # Special handling for stream parameter (it's a Proc)
+                        metadata[field] = if field == "stream"
+                          true  # Just mark as streaming
+                        else
+                          parameters[field_sym]
+                        end
+                      end
                     end
-                  end
-                end
-
-                # Set input messages as JSON
-                if parameters[:messages]
-                  span.set_attribute("braintrust.input_json", JSON.generate(parameters[:messages]))
-                end
-
-                begin
-                  # Call the original method
-                  response = super(parameters: parameters)
 
-                  # Set output (choices) as JSON
-                  if response && response["choices"]&.any?
-                    span.set_attribute("braintrust.output_json", JSON.generate(response["choices"]))
-                  end
-
-                  # Set metrics (token usage)
-                  if response && response["usage"]
-                    metrics = Braintrust::Trace::AlexRudall::RubyOpenAI.parse_usage_tokens(response["usage"])
-                    span.set_attribute("braintrust.metrics", JSON.generate(metrics)) unless metrics.empty?
-                  end
+                    # Set input messages as JSON
+                    if parameters[:messages]
+                      span.set_attribute("braintrust.input_json", JSON.generate(parameters[:messages]))
+                    end
 
-                  # Add response metadata fields
-                  if response
-                    metadata["id"] = response["id"] if response["id"]
-                    metadata["created"] = response["created"] if response["created"]
-                    metadata["system_fingerprint"] = response["system_fingerprint"] if response["system_fingerprint"]
-                    metadata["service_tier"] = response["service_tier"] if response["service_tier"]
+                    begin
+                      # Call the original method
+                      response = super(parameters: parameters)
+
+                      # Set output (choices) as JSON
+                      if response && response["choices"]&.any?
+                        span.set_attribute("braintrust.output_json", JSON.generate(response["choices"]))
+                      end
+
+                      # Set metrics (token usage)
+                      if response && response["usage"]
+                        metrics = Braintrust::Trace::Contrib::Github::Alexrudall::RubyOpenAI.parse_usage_tokens(response["usage"])
+                        span.set_attribute("braintrust.metrics", JSON.generate(metrics)) unless metrics.empty?
+                      end
+
+                      # Add response metadata fields
+                      if response
+                        metadata["id"] = response["id"] if response["id"]
+                        metadata["created"] = response["created"] if response["created"]
+                        metadata["system_fingerprint"] = response["system_fingerprint"] if response["system_fingerprint"]
+                        metadata["service_tier"] = response["service_tier"] if response["service_tier"]
+                      end
+
+                      # Set metadata ONCE at the end with complete hash
+                      span.set_attribute("braintrust.metadata", JSON.generate(metadata))
+
+                      response
+                    rescue => e
+                      # Record exception in span
+                      span.record_exception(e)
+                      span.status = OpenTelemetry::Trace::Status.error("Exception: #{e.class} - #{e.message}")
+                      raise
+                    end
                   end
-
-                  # Set metadata ONCE at the end with complete hash
-                  span.set_attribute("braintrust.metadata", JSON.generate(metadata))
-
-                  response
-                rescue => e
-                  # Record exception in span
-                  span.record_exception(e)
-                  span.status = OpenTelemetry::Trace::Status.error("Exception: #{e.class} - #{e.message}")
-                  raise
                 end
               end
+
+              # Prepend the wrapper to the client's singleton class
+              client.singleton_class.prepend(wrapper)
             end
           end
-
-          # Prepend the wrapper to the client's singleton class
-          client.singleton_class.prepend(wrapper)
         end
       end
     end
+
+    # Backwards compatibility: this module was originally at Braintrust::Trace::AlexRudall::RubyOpenAI
+    module AlexRudall
+      RubyOpenAI = Contrib::Github::Alexrudall::RubyOpenAI
+    end
   end
 end

data/lib/braintrust/trace/contrib/github.com/crmne/ruby_llm.rb

@@ -0,0 +1,447 @@
+# frozen_string_literal: true
+
+require "opentelemetry/sdk"
+require "json"
+require_relative "../../../tokens"
+require_relative "../../../../logger"
+
+module Braintrust
+  module Trace
+    module Contrib
+      module Github
+        module Crmne
+          module RubyLLM
+            # Helper to safely set a JSON attribute on a span
+            # Only sets the attribute if obj is present
+            # @param span [OpenTelemetry::Trace::Span] the span to set attribute on
+            # @param attr_name [String] the attribute name (e.g., "braintrust.output_json")
+            # @param obj [Object] the object to serialize to JSON
+            # @return [void]
+            def self.set_json_attr(span, attr_name, obj)
+              return unless obj
+              span.set_attribute(attr_name, JSON.generate(obj))
+            rescue => e
+              Log.debug("Failed to serialize #{attr_name}: #{e.message}")
+            end
+
+            # Parse usage tokens from RubyLLM response
+            # RubyLLM uses Anthropic-style field naming (input_tokens, output_tokens)
+            # @param usage [Hash, Object] usage object from RubyLLM response
+            # @return [Hash<String, Integer>] metrics hash with normalized names
+            def self.parse_usage_tokens(usage)
+              Braintrust::Trace.parse_anthropic_usage_tokens(usage)
+            end
+
+            # Wrap RubyLLM to automatically create spans for chat requests
+            # Supports both synchronous and streaming requests
+            #
+            # Usage:
+            #   # Wrap the class once (affects all future instances):
+            #   Braintrust::Trace::Contrib::Github::Crmne::RubyLLM.wrap
+            #
+            #   # Or wrap a specific instance:
+            #   chat = RubyLLM.chat(model: "gpt-4o-mini")
+            #   Braintrust::Trace::Contrib::Github::Crmne::RubyLLM.wrap(chat)
+            #
+            # @param chat [RubyLLM::Chat, nil] the RubyLLM chat instance to wrap (if nil, wraps the class)
+            # @param tracer_provider [OpenTelemetry::SDK::Trace::TracerProvider] the tracer provider (defaults to global)
+            def self.wrap(chat = nil, tracer_provider: nil)
+              tracer_provider ||= ::OpenTelemetry.tracer_provider
+
+              # If no chat instance provided, wrap the class globally via initialize hook
+              if chat.nil?
+                return if defined?(::RubyLLM::Chat) && ::RubyLLM::Chat.instance_variable_defined?(:@braintrust_wrapper_module)
+
+                # Create module that wraps initialize to auto-wrap each new instance
+                wrapper_module = Module.new do
+                  define_method(:initialize) do |*args, **kwargs, &block|
+                    super(*args, **kwargs, &block)
+                    # Auto-wrap this instance during initialization
+                    Braintrust::Trace::Contrib::Github::Crmne::RubyLLM.wrap(self, tracer_provider: tracer_provider)
+                    self
+                  end
+                end
+
+                # Store reference to wrapper module for cleanup
+                ::RubyLLM::Chat.instance_variable_set(:@braintrust_wrapper_module, wrapper_module)
+                ::RubyLLM::Chat.prepend(wrapper_module)
+                return nil
+              end
+
+              # Check if already wrapped to make this idempotent
+              return chat if chat.instance_variable_get(:@braintrust_wrapped)
+
+              # Create a wrapper module that intercepts chat.ask
+              wrapper = create_wrapper_module(tracer_provider)
+
+              # Mark as wrapped and prepend the wrapper to the chat instance
+              chat.instance_variable_set(:@braintrust_wrapped, true)
+              chat.singleton_class.prepend(wrapper)
+              chat
+            end
+
+            # Unwrap RubyLLM to remove Braintrust tracing
+            # For class-level unwrapping, removes the initialize override from the wrapper module
+            # For instance-level unwrapping, clears the wrapped flag
+            #
+            # @param chat [RubyLLM::Chat, nil] the RubyLLM chat instance to unwrap (if nil, unwraps the class)
+            def self.unwrap(chat = nil)
+              # If no chat instance provided, unwrap the class globally
+              if chat.nil?
+                if defined?(::RubyLLM::Chat) && ::RubyLLM::Chat.instance_variable_defined?(:@braintrust_wrapper_module)
+                  wrapper_module = ::RubyLLM::Chat.instance_variable_get(:@braintrust_wrapper_module)
+                  # Redefine initialize to just call super (disables auto-wrapping)
+                  # We can't actually remove a prepended module, so we make it a no-op
+                  wrapper_module.module_eval do
+                    define_method(:initialize) do |*args, **kwargs, &block|
+                      super(*args, **kwargs, &block)
+                    end
+                  end
+                  ::RubyLLM::Chat.remove_instance_variable(:@braintrust_wrapper_module)
+                end
+                return nil
+              end
+
+              # Unwrap instance
+              chat.remove_instance_variable(:@braintrust_wrapped) if chat.instance_variable_defined?(:@braintrust_wrapped)
+              chat
+            end
+
+            # Wrap the RubyLLM::Chat class globally
+            # @param tracer_provider [OpenTelemetry::SDK::Trace::TracerProvider] the tracer provider
+            def self.wrap_class(tracer_provider)
+              return unless defined?(::RubyLLM::Chat)
+
+              wrapper = create_wrapper_module(tracer_provider)
+              ::RubyLLM::Chat.prepend(wrapper)
+            end
+
+            # Create the wrapper module that intercepts chat.ask
+            # @param tracer_provider [OpenTelemetry::SDK::Trace::TracerProvider] the tracer provider
+            # @return [Module] the wrapper module
+            def self.create_wrapper_module(tracer_provider)
+              Module.new do
+                define_method(:ask) do |prompt = nil, **params, &block|
+                  tracer = tracer_provider.tracer("braintrust")
+
+                  if block
+                    # Handle streaming request
+                    wrapped_block = proc do |chunk|
+                      block.call(chunk)
+                    end
+                    Braintrust::Trace::Contrib::Github::Crmne::RubyLLM.handle_streaming_ask(self, tracer, prompt, params, block) do |aggregated_chunks|
+                      super(prompt, **params) do |chunk|
+                        aggregated_chunks << chunk
+                        wrapped_block.call(chunk)
+                      end
+                    end
+                  else
+                    # Handle non-streaming request
+                    Braintrust::Trace::Contrib::Github::Crmne::RubyLLM.handle_non_streaming_ask(self, tracer, prompt, params) do
+                      super(prompt, **params)
+                    end
+                  end
+                end
+              end
+            end
+
+            # Handle streaming chat request with tracing
+            # @param chat [RubyLLM::Chat] the chat instance
+            # @param tracer [OpenTelemetry::Trace::Tracer] the tracer
+            # @param prompt [String, nil] the user prompt
+            # @param params [Hash] additional parameters
+            # @param block [Proc] the streaming block
+            def self.handle_streaming_ask(chat, tracer, prompt, params, block)
+              # Start span immediately for accurate timing
+              span = tracer.start_span("ruby_llm.chat.ask")
+
+              aggregated_chunks = []
+
+              # Extract metadata and build input messages
+              metadata = extract_metadata(chat, stream: true)
+              input_messages = build_input_messages(chat, prompt)
+
+              # Set input and metadata
+              set_json_attr(span, "braintrust.input_json", input_messages) if input_messages.any?
+              set_json_attr(span, "braintrust.metadata", metadata)
+
+              # Call original method, passing aggregated_chunks to the block
+              begin
+                result = yield aggregated_chunks
+              rescue => e
+                span.record_exception(e)
+                span.status = ::OpenTelemetry::Trace::Status.error("RubyLLM error: #{e.message}")
+                span.finish
+                raise
+              end
+
+              # Set output and metrics from aggregated chunks
+              capture_streaming_output(span, aggregated_chunks, result)
+              span.finish
+              result
+            end
+
+            # Handle non-streaming chat request with tracing
+            # @param chat [RubyLLM::Chat] the chat instance
+            # @param tracer [OpenTelemetry::Trace::Tracer] the tracer
+            # @param prompt [String, nil] the user prompt
+            # @param params [Hash] additional parameters
+            def self.handle_non_streaming_ask(chat, tracer, prompt, params)
+              # Start span immediately for accurate timing
+              span = tracer.start_span("ruby_llm.chat.ask")
+
+              begin
+                # Extract metadata and build input messages
+                metadata = extract_metadata(chat)
+                input_messages = build_input_messages(chat, prompt)
+                set_json_attr(span, "braintrust.input_json", input_messages) if input_messages.any?
+
+                # Remember message count before the call (for tool call detection)
+                messages_before_count = (chat.respond_to?(:messages) && chat.messages) ? chat.messages.length : 0
+
+                # Call the original method
+                response = yield
+
+                # Capture output and metrics
+                capture_non_streaming_output(span, chat, response, messages_before_count)
+
+                # Set metadata
+                set_json_attr(span, "braintrust.metadata", metadata)
+
+                response
+              ensure
+                span.finish
+              end
+            end
+
+            # Extract metadata from chat instance (provider, model, tools, stream flag)
+            # @param chat [RubyLLM::Chat] the chat instance
+            # @param stream [Boolean] whether this is a streaming request
+            # @return [Hash] metadata hash
+            def self.extract_metadata(chat, stream: false)
+              metadata = {"provider" => "ruby_llm"}
+              metadata["stream"] = true if stream
+
+              # Extract model
+              if chat.respond_to?(:model) && chat.model
+                model = chat.model.respond_to?(:id) ? chat.model.id : chat.model.to_s
+                metadata["model"] = model
+              end
+
+              # Extract tools (only for non-streaming)
+              if !stream && chat.respond_to?(:tools) && chat.tools&.any?
+                metadata["tools"] = extract_tools_metadata(chat)
+              end
+
+              metadata
+            end
+
+            # Extract tools metadata from chat instance
+            # @param chat [RubyLLM::Chat] the chat instance
+            # @return [Array<Hash>] array of tool schemas
+            def self.extract_tools_metadata(chat)
+              provider = chat.instance_variable_get(:@provider) if chat.instance_variable_defined?(:@provider)
+
+              chat.tools.map do |_name, tool|
+                format_tool_schema(tool, provider)
+              end
+            end
+
+            # Format a tool into OpenAI-compatible schema
+            # @param tool [Object] the tool object
+            # @param provider [Object, nil] the provider instance
+            # @return [Hash] tool schema
+            def self.format_tool_schema(tool, provider)
+              tool_schema = nil
+
+              # Use provider-specific tool_for method if available
+              if provider
+                begin
+                  tool_schema = if provider.is_a?(::RubyLLM::Providers::OpenAI)
+                    ::RubyLLM::Providers::OpenAI::Tools.tool_for(tool)
+                  elsif defined?(::RubyLLM::Providers::Anthropic) && provider.is_a?(::RubyLLM::Providers::Anthropic)
+                    ::RubyLLM::Providers::Anthropic::Tools.tool_for(tool)
+                  elsif tool.respond_to?(:params_schema) && tool.params_schema
+                    build_basic_tool_schema(tool)
+                  else
+                    build_minimal_tool_schema(tool)
+                  end
+                rescue NameError, ArgumentError => e
+                  # If provider-specific tool_for fails, fall back to basic format
+                  Log.debug("Failed to extract tool schema using provider-specific method: #{e.class.name}: #{e.message}")
+                  tool_schema = (tool.respond_to?(:params_schema) && tool.params_schema) ? build_basic_tool_schema(tool) : build_minimal_tool_schema(tool)
+                end
+              else
+                # No provider, use basic format with params_schema if available
+                tool_schema = (tool.respond_to?(:params_schema) && tool.params_schema) ? build_basic_tool_schema(tool) : build_minimal_tool_schema(tool)
+              end
+
+              # Strip RubyLLM-specific fields to match native OpenAI format
+              # Handle both symbol and string keys
+              function_key = tool_schema&.key?(:function) ? :function : "function"
+              if tool_schema && tool_schema[function_key]
+                tool_params = tool_schema[function_key][:parameters] || tool_schema[function_key]["parameters"]
+                if tool_params.is_a?(Hash)
+                  tool_params.delete("strict")
+                  tool_params.delete(:strict)
+                  tool_params.delete("additionalProperties")
+                  tool_params.delete(:additionalProperties)
+                end
+              end
+
+              tool_schema
+            end
+
+            # Build a basic tool schema with parameters
+            # @param tool [Object] the tool object
+            # @return [Hash] tool schema
+            def self.build_basic_tool_schema(tool)
+              {
+                "type" => "function",
+                "function" => {
+                  "name" => tool.name.to_s,
+                  "description" => tool.description,
+                  "parameters" => tool.params_schema
+                }
+              }
+            end
+
+            # Build a minimal tool schema without parameters
+            # @param tool [Object] the tool object
+            # @return [Hash] tool schema
+            def self.build_minimal_tool_schema(tool)
+              {
+                "type" => "function",
+                "function" => {
+                  "name" => tool.name.to_s,
+                  "description" => tool.description,
+                  "parameters" => {}
+                }
+              }
+            end
+
+            # Build input messages array from chat history and prompt
+            # @param chat [RubyLLM::Chat] the chat instance
+            # @param prompt [String, nil] the user prompt
+            # @return [Array<Hash>] array of message hashes
+            def self.build_input_messages(chat, prompt)
+              input_messages = []
+
+              # Add conversation history
+              if chat.respond_to?(:messages) && chat.messages&.any?
+                input_messages = chat.messages.map { |m| m.respond_to?(:to_h) ? m.to_h : m }
+              end
+
+              # Add current prompt
+              input_messages << {role: "user", content: prompt} if prompt
+
+              input_messages
+            end
+
+            # Capture streaming output and metrics
+            # @param span [OpenTelemetry::Trace::Span] the span
+            # @param aggregated_chunks [Array] the aggregated chunks
+            # @param result [Object] the result object
+            def self.capture_streaming_output(span, aggregated_chunks, result)
+              return if aggregated_chunks.empty?
+
+              # Aggregate content from chunks
+              aggregated_content = aggregated_chunks.map { |c|
+                c.respond_to?(:content) ? c.content : c.to_s
+              }.join
+
+              output = [{
+                role: "assistant",
+                content: aggregated_content
+              }]
+              set_json_attr(span, "braintrust.output_json", output)
+
+              # Try to extract usage from the result
+              if result.respond_to?(:usage) && result.usage
+                metrics = parse_usage_tokens(result.usage)
+                set_json_attr(span, "braintrust.metrics", metrics) unless metrics.empty?
+              end
+            end
+
+            # Capture non-streaming output and metrics
+            # @param span [OpenTelemetry::Trace::Span] the span
+            # @param chat [RubyLLM::Chat] the chat instance
+            # @param response [Object] the response object
+            # @param messages_before_count [Integer] message count before the call
+            def self.capture_non_streaming_output(span, chat, response, messages_before_count)
+              return unless response
+
+              # Build message object from response
+              message = {
+                "role" => "assistant",
+                "content" => nil
+              }
+
+              # Add content if it's a simple text response
+              if response.respond_to?(:content) && response.content && !response.content.empty?
+                message["content"] = response.content
+              end
+
+              # Check if there are tool calls in the messages history
+              if chat.respond_to?(:messages) && chat.messages
+                assistant_msg = chat.messages[(messages_before_count + 1)..].find { |m|
+                  m.role.to_s == "assistant" && m.respond_to?(:tool_calls) && m.tool_calls&.any?
+                }
+
+                if assistant_msg&.tool_calls&.any?
+                  message["tool_calls"] = format_tool_calls(assistant_msg.tool_calls)
+                  message["content"] = nil
+                end
+              end
+
+              # Format as OpenAI choices[] structure
+              output = [{
+                "index" => 0,
+                "message" => message,
+                "finish_reason" => message["tool_calls"] ? "tool_calls" : "stop"
+              }]
+
+              set_json_attr(span, "braintrust.output_json", output)
+
+              # Set metrics (token usage)
+              if response.respond_to?(:to_h)
+                response_hash = response.to_h
+                usage = {
+                  "input_tokens" => response_hash[:input_tokens],
+                  "output_tokens" => response_hash[:output_tokens],
+                  "cached_tokens" => response_hash[:cached_tokens],
+                  "cache_creation_tokens" => response_hash[:cache_creation_tokens]
+                }.compact
+
+                unless usage.empty?
+                  metrics = parse_usage_tokens(usage)
+                  set_json_attr(span, "braintrust.metrics", metrics) unless metrics.empty?
+                end
+              end
+            end
+
+            # Format tool calls into OpenAI format
+            # @param tool_calls [Hash, Array] the tool calls
+            # @return [Array<Hash>] formatted tool calls
+            def self.format_tool_calls(tool_calls)
+              tool_calls.map do |_id, tc|
+                # Ensure arguments is a JSON string (OpenAI format)
+                args = tc.arguments
+                args_string = args.is_a?(String) ? args : JSON.generate(args)
+
+                {
+                  "id" => tc.id,
+                  "type" => "function",
+                  "function" => {
+                    "name" => tc.name,
+                    "arguments" => args_string
+                  }
+                }
+              end
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/braintrust/trace/contrib/openai.rb

@@ -2,6 +2,7 @@
 
 require "opentelemetry/sdk"
 require "json"
+require_relative "../tokens"
 
 module Braintrust
   module Trace
@@ -17,72 +18,11 @@ module Braintrust
         span.set_attribute(attr_name, JSON.generate(obj))
       end
 
-      # Parse usage tokens from OpenAI API response, handling nested token_details
-      # Maps OpenAI field names to Braintrust standard names:
-      # - input_tokens → prompt_tokens
-      # - output_tokens → completion_tokens
-      # - total_tokens → tokens
-      # - *_tokens_details.* → prefix_*
-      #
+      # Parse usage tokens from OpenAI API response
       # @param usage [Hash, Object] usage object from OpenAI response
       # @return [Hash<String, Integer>] metrics hash with normalized names
       def self.parse_usage_tokens(usage)
-        metrics = {}
-        return metrics unless usage
-
-        # Convert to hash if it's an object
-        usage_hash = usage.respond_to?(:to_h) ? usage.to_h : usage
-
-        usage_hash.each do |key, value|
-          key_str = key.to_s
-
-          # Handle nested *_tokens_details objects
-          if key_str.end_with?("_tokens_details")
-            # Convert to hash if it's an object (OpenAI gem returns objects)
-            details_hash = value.respond_to?(:to_h) ? value.to_h : value
-            next unless details_hash.is_a?(Hash)
-
-            # Extract prefix (e.g., "prompt" from "prompt_tokens_details")
-            prefix = key_str.sub(/_tokens_details$/, "")
-            # Translate "input" → "prompt", "output" → "completion"
-            prefix = translate_metric_prefix(prefix)
-
-            # Process nested fields (e.g., cached_tokens, reasoning_tokens)
-            details_hash.each do |detail_key, detail_value|
-              next unless detail_value.is_a?(Numeric)
-              metrics["#{prefix}_#{detail_key}"] = detail_value.to_i
-            end
-          elsif value.is_a?(Numeric)
-            # Handle top-level token fields
-            case key_str
-            when "input_tokens"
-              metrics["prompt_tokens"] = value.to_i
-            when "output_tokens"
-              metrics["completion_tokens"] = value.to_i
-            when "total_tokens"
-              metrics["tokens"] = value.to_i
-            else
-              # Keep other numeric fields as-is (future-proofing)
-              metrics[key_str] = value.to_i
-            end
-          end
-        end
-
-        metrics
-      end
-
-      # Translate metric prefix to be consistent between different API formats
-      # @param prefix [String] the prefix to translate
-      # @return [String] translated prefix
-      def self.translate_metric_prefix(prefix)
-        case prefix
-        when "input"
-          "prompt"
-        when "output"
-          "completion"
-        else
-          prefix
-        end
+        Braintrust::Trace.parse_openai_usage_tokens(usage)
       end
 
       # Aggregate streaming chunks into a single response structure

data/lib/braintrust/trace/tokens.rb

@@ -0,0 +1,101 @@
+# frozen_string_literal: true
+
+module Braintrust
+  module Trace
+    # Parse OpenAI usage tokens into normalized Braintrust metrics.
+    # Handles standard fields and *_tokens_details nested objects.
+    # @param usage [Hash, Object] usage object from OpenAI response
+    # @return [Hash<String, Integer>] normalized metrics
+    def self.parse_openai_usage_tokens(usage)
+      metrics = {}
+      return metrics unless usage
+
+      usage_hash = usage.respond_to?(:to_h) ? usage.to_h : usage
+      return metrics unless usage_hash.is_a?(Hash)
+
+      # Field mappings: OpenAI → Braintrust
+      field_map = {
+        "prompt_tokens" => "prompt_tokens",
+        "completion_tokens" => "completion_tokens",
+        "total_tokens" => "tokens"
+      }
+
+      # Prefix mappings for *_tokens_details
+      prefix_map = {
+        "prompt" => "prompt",
+        "completion" => "completion"
+      }
+
+      usage_hash.each do |key, value|
+        key_str = key.to_s
+
+        if value.is_a?(Numeric)
+          target = field_map[key_str]
+          metrics[target] = value.to_i if target
+        elsif key_str.end_with?("_tokens_details")
+          # Convert to hash if it's an object (OpenAI SDK returns objects)
+          details_hash = value.respond_to?(:to_h) ? value.to_h : value
+          next unless details_hash.is_a?(Hash)
+
+          raw_prefix = key_str.sub(/_tokens_details$/, "")
+          prefix = prefix_map[raw_prefix] || raw_prefix
+          details_hash.each do |detail_key, detail_value|
+            next unless detail_value.is_a?(Numeric)
+            metrics["#{prefix}_#{detail_key}"] = detail_value.to_i
+          end
+        end
+      end
+
+      # Calculate total if missing
+      if !metrics.key?("tokens") && metrics.key?("prompt_tokens") && metrics.key?("completion_tokens")
+        metrics["tokens"] = metrics["prompt_tokens"] + metrics["completion_tokens"]
+      end
+
+      metrics
+    end
+
+    # Parse Anthropic usage tokens into normalized Braintrust metrics.
+    # Accumulates cache tokens into prompt_tokens and calculates total.
+    # @param usage [Hash, Object] usage object from Anthropic response
+    # @return [Hash<String, Integer>] normalized metrics
+    def self.parse_anthropic_usage_tokens(usage)
+      metrics = {}
+      return metrics unless usage
+
+      usage_hash = usage.respond_to?(:to_h) ? usage.to_h : usage
+      return metrics unless usage_hash.is_a?(Hash)
+
+      # Field mappings: Anthropic → Braintrust
+      # Also handles RubyLLM's simplified cache field names
+      field_map = {
+        "input_tokens" => "prompt_tokens",
+        "output_tokens" => "completion_tokens",
+        "cache_read_input_tokens" => "prompt_cached_tokens",
+        "cache_creation_input_tokens" => "prompt_cache_creation_tokens",
+        # RubyLLM uses simplified names
+        "cached_tokens" => "prompt_cached_tokens",
+        "cache_creation_tokens" => "prompt_cache_creation_tokens"
+      }
+
+      usage_hash.each do |key, value|
+        next unless value.is_a?(Numeric)
+        key_str = key.to_s
+        target = field_map[key_str]
+        metrics[target] = value.to_i if target
+      end
+
+      # Accumulate cache tokens into prompt_tokens (matching TS/Python SDKs)
+      prompt_tokens = (metrics["prompt_tokens"] || 0) +
+        (metrics["prompt_cached_tokens"] || 0) +
+        (metrics["prompt_cache_creation_tokens"] || 0)
+      metrics["prompt_tokens"] = prompt_tokens if prompt_tokens > 0
+
+      # Calculate total
+      if metrics.key?("prompt_tokens") && metrics.key?("completion_tokens")
+        metrics["tokens"] = metrics["prompt_tokens"] + metrics["completion_tokens"]
+      end
+
+      metrics
+    end
+  end
+end

data/lib/braintrust/trace.rb

@@ -40,6 +40,23 @@ rescue LoadError
   # Anthropic gem not installed - integration will not be available
 end
 
+# RubyLLM integration is optional - automatically loaded if ruby_llm gem is available
+#
+# Usage:
+#   # Wrap the class once (affects all instances):
+#   Braintrust::Trace::RubyLLM.wrap
+#
+#   # Or wrap a specific instance:
+#   chat = RubyLLM.chat(model: "gpt-4o-mini")
+#   Braintrust::Trace::RubyLLM.wrap(chat)
+#
+begin
+  require "ruby_llm"
+  require_relative "trace/contrib/github.com/crmne/ruby_llm"
+rescue LoadError
+  # RubyLLM gem not installed - integration will not be available
+end
+
 module Braintrust
   module Trace
     # Set up OpenTelemetry tracing with Braintrust

data/lib/braintrust/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module Braintrust
-  VERSION = "0.0.5"
+  VERSION = "0.0.6"
 end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: braintrust
 version: !ruby/object:Gem::Version
-  version: 0.0.5
+  version: 0.0.6
 platform: ruby
 authors:
 - Braintrust
@@ -205,9 +205,11 @@ files:
 - lib/braintrust/trace/attachment.rb
 - lib/braintrust/trace/contrib/anthropic.rb
 - lib/braintrust/trace/contrib/github.com/alexrudall/ruby-openai/ruby-openai.rb
+- lib/braintrust/trace/contrib/github.com/crmne/ruby_llm.rb
 - lib/braintrust/trace/contrib/openai.rb
 - lib/braintrust/trace/span_filter.rb
 - lib/braintrust/trace/span_processor.rb
+- lib/braintrust/trace/tokens.rb
 - lib/braintrust/version.rb
 homepage: https://github.com/braintrustdata/braintrust-sdk-ruby
 licenses: