braintrust 0.0.4 → 0.0.6

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 39d85e02bd85a931ee7f16de103d48d1184048e3ad8d791eda37bc323a653716
-  data.tar.gz: a0b1d5493e8ad3004007e78d608154077a33c92a436bce23eb36cfbe94c3bdd4
+  metadata.gz: 866cb2e797502f00cda1625ad90f4d734b4b83f0d21d8243675a933fae9df693
+  data.tar.gz: f74151b0e18b12cf19b61b1b75b2f58e784d4171f21c0996526d29c719174260
 SHA512:
-  metadata.gz: a5dcbd1b2bf2c0ab2355ff36c9cfce4fe10e175c0aa8df80ea3176be4002271744ca3a9fd7ef52cec888e0b326772518554921f7d657a79ba347b26c4c93b80c
-  data.tar.gz: 78677bd57e6ed1778f74b87e050dd5bbfdc8390e73f919aa57ea680cd2cd4338086e5df4982274c3fd62e690d34ec81078d7c76e75a467ab2f5b0667e6d530d6
+  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
@@ -196,7 +147,6 @@ module Braintrust
         wrapper = Module.new do
           define_method(:stream) do |**params, &block|
             tracer = tracer_provider.tracer("braintrust")
-            aggregated_events = []
 
             metadata = {
               "provider" => "anthropic",
@@ -256,183 +206,110 @@ module Braintrust
             end
 
             # Store references on the stream object itself for the wrapper
-            stream.instance_variable_set(:@braintrust_aggregated_events, aggregated_events)
             stream.instance_variable_set(:@braintrust_span, span)
             stream.instance_variable_set(:@braintrust_metadata, metadata)
+            stream.instance_variable_set(:@braintrust_span_finished, false)
 
-            # Wrap the stream to aggregate events
-            original_each = stream.method(:each)
-            stream.define_singleton_method(:each) do |&user_block|
-              events = instance_variable_get(:@braintrust_aggregated_events)
-              span_obj = instance_variable_get(:@braintrust_span)
-              meta = instance_variable_get(:@braintrust_metadata)
-
-              begin
-                original_each.call do |event|
-                  # Store event data for aggregation
-                  events << event.to_h if event.respond_to?(:to_h)
-                  # Call user's block if provided
-                  user_block&.call(event)
-                end
-              rescue => e
-                span_obj.record_exception(e)
-                span_obj.status = ::OpenTelemetry::Trace::Status.error("Streaming error: #{e.message}")
-                raise
-              ensure
-                # Always aggregate and finish span after stream completes
-                unless events.empty?
-                  aggregated_output = Braintrust::Trace::Anthropic.aggregate_streaming_events(events)
-
-                  # Set output
-                  if aggregated_output[:content]
-                    output = [{
-                      role: "assistant",
-                      content: aggregated_output[:content]
-                    }]
-                    Braintrust::Trace::Anthropic.set_json_attr(span_obj, "braintrust.output_json", output)
-                  end
-
-                  # Set metrics if usage is available
-                  if aggregated_output[:usage]
-                    metrics = Braintrust::Trace::Anthropic.parse_usage_tokens(aggregated_output[:usage])
-                    Braintrust::Trace::Anthropic.set_json_attr(span_obj, "braintrust.metrics", metrics) unless metrics.empty?
-                  end
-
-                  # Update metadata with response fields
-                  meta["stop_reason"] = aggregated_output[:stop_reason] if aggregated_output[:stop_reason]
-                  meta["model"] = aggregated_output[:model] if aggregated_output[:model]
-                  Braintrust::Trace::Anthropic.set_json_attr(span_obj, "braintrust.metadata", meta)
-                end
+            # Local helper for brevity
+            set_json_attr = ->(attr_name, obj) { Braintrust::Trace::Anthropic.set_json_attr(span, attr_name, obj) }
+
+            # Helper lambda to extract stream data and set span attributes
+            # This is DRY - used by both .each() and .text() wrappers
+            extract_stream_metadata = lambda do
+              # Extract the SDK's internal accumulated message (built during streaming)
+              acc_msg = stream.instance_variable_get(:@accumated_message_snapshot)
+              return unless acc_msg
 
-                span_obj.finish
+              # Set output from accumulated message
+              if acc_msg.respond_to?(:content) && acc_msg.content
+                content_array = acc_msg.content.map(&:to_h)
+                output = [{
+                  role: acc_msg.respond_to?(:role) ? acc_msg.role : "assistant",
+                  content: content_array
+                }]
+                set_json_attr.call("braintrust.output_json", output)
               end
-            end
 
-            # If a block was provided to stream(), call each with it immediately
-            if block
-              stream.each(&block)
+              # Set metrics from accumulated message
+              if acc_msg.respond_to?(:usage) && acc_msg.usage
+                metrics = Braintrust::Trace::Anthropic.parse_usage_tokens(acc_msg.usage)
+                set_json_attr.call("braintrust.metrics", metrics) unless metrics.empty?
+              end
+
+              # Update metadata with response fields
+              if acc_msg.respond_to?(:stop_reason) && acc_msg.stop_reason
+                metadata["stop_reason"] = acc_msg.stop_reason
+              end
+              if acc_msg.respond_to?(:model) && acc_msg.model
+                metadata["model"] = acc_msg.model
+              end
+              set_json_attr.call("braintrust.metadata", metadata)
             end
 
-            stream
-          end
-        end
+            # Helper lambda to finish span (prevents double-finishing via closure)
+            finish_braintrust_span = lambda do
+              return if stream.instance_variable_get(:@braintrust_span_finished)
+              stream.instance_variable_set(:@braintrust_span_finished, true)
 
-        # Prepend the wrapper to the messages resource
-        client.messages.singleton_class.prepend(wrapper)
-      end
+              extract_stream_metadata.call
+              span.finish
+            end
 
-      # Aggregate streaming events into a single response structure
-      # @param events [Array<Hash>] array of event hashes from stream
-      # @return [Hash] aggregated response with content, usage, etc.
-      def self.aggregate_streaming_events(events)
-        return {} if events.empty?
-
-        result = {
-          content: [],
-          usage: {},
-          stop_reason: nil,
-          model: nil
-        }
-
-        # Track content blocks by index
-        content_blocks = {}
-        content_builders = {}
-
-        events.each do |event|
-          event_type = event[:type] || event["type"]
-          next unless event_type
-
-          case event_type
-          when "message_start"
-            # Extract model and initial usage (input tokens, cache tokens)
-            message = event[:message] || event["message"]
-            if message
-              result[:model] = message[:model] || message["model"]
-              if message[:usage] || message["usage"]
-                usage = message[:usage] || message["usage"]
-                result[:usage].merge!(usage)
-              end
+            # Wrap .each() to ensure span finishes after consumption
+            original_each = stream.method(:each)
+            stream.define_singleton_method(:each) do |&user_block|
+              # Consume stream, calling user's block for each event
+              # The SDK builds @accumated_message_snapshot internally
+              original_each.call(&user_block)
+            rescue => e
+              span.record_exception(e)
+              span.status = ::OpenTelemetry::Trace::Status.error("Streaming error: #{e.message}")
+              raise
+            ensure
+              # Extract accumulated message and finish span
+              finish_braintrust_span.call
             end
 
-          when "content_block_start"
-            # Initialize a new content block
-            index = event[:index] || event["index"]
-            content_block = event[:content_block] || event["content_block"]
-            content_blocks[index] = content_block if index && content_block
-
-          when "content_block_delta"
-            # Accumulate deltas for content blocks
-            index = event[:index] || event["index"]
-            delta = event[:delta] || event["delta"]
-            next unless index && delta
-
-            delta_type = delta[:type] || delta["type"]
-            content_blocks[index] ||= {}
-
-            case delta_type
-            when "text_delta"
-              # Accumulate text
-              text = delta[:text] || delta["text"]
-              if text
-                content_builders[index] ||= ""
-                content_builders[index] += text
-                content_blocks[index][:type] = "text"
-                content_blocks[index]["type"] = "text"
-              end
+            # Wrap .text() to return an Enumerable that ensures span finishes
+            original_text = stream.method(:text)
+            stream.define_singleton_method(:text) do
+              text_enum = original_text.call
 
-            when "input_json_delta"
-              # Accumulate JSON for tool_use blocks
-              partial_json = delta[:partial_json] || delta["partial_json"]
-              if partial_json
-                content_builders[index] ||= ""
-                content_builders[index] += partial_json
-                content_blocks[index][:type] = "tool_use"
-                content_blocks[index]["type"] = "tool_use"
+              # Return wrapper Enumerable that finishes span after consumption
+              Enumerator.new do |y|
+                # Consume text enumerable (this consumes underlying stream)
+                # The SDK builds @accumated_message_snapshot internally
+                text_enum.each { |text| y << text }
+              rescue => e
+                span.record_exception(e)
+                span.status = ::OpenTelemetry::Trace::Status.error("Streaming error: #{e.message}")
+                raise
+              ensure
+                # Extract accumulated message and finish span
+                finish_braintrust_span.call
               end
             end
 
-          when "message_delta"
-            # Get final stop reason and cumulative usage (output tokens)
-            delta = event[:delta] || event["delta"]
-            if delta
-              stop_reason = delta[:stop_reason] || delta["stop_reason"]
-              result[:stop_reason] = stop_reason if stop_reason
+            # Wrap .close() to ensure span finishes even if stream not consumed
+            original_close = stream.method(:close)
+            stream.define_singleton_method(:close) do
+              original_close.call
+            ensure
+              # Finish span even if stream was closed early
+              finish_braintrust_span.call
             end
 
-            usage = event[:usage] || event["usage"]
-            result[:usage].merge!(usage) if usage
-          end
-        end
-
-        # Build final content array from aggregated blocks
-        content_builders.each do |index, text|
-          block = content_blocks[index]
-          next unless block
-
-          block_type = block[:type] || block["type"]
-          case block_type
-          when "text"
-            block[:text] = text
-            block["text"] = text
-          when "tool_use"
-            # Parse the accumulated JSON string
-            begin
-              parsed = JSON.parse(text)
-              block[:input] = parsed
-              block["input"] = parsed
-            rescue JSON::ParserError
-              block[:input] = text
-              block["input"] = text
+            # If a block was provided to stream(), call each with it immediately
+            if block
+              stream.each(&block)
             end
-          end
-        end
 
-        # Convert blocks hash to sorted array
-        if content_blocks.any?
-          result[:content] = content_blocks.keys.sort.map { |idx| content_blocks[idx] }
+            stream
+          end
         end
 
-        result
+        # Prepend the wrapper to the messages resource
+        client.messages.singleton_class.prepend(wrapper)
       end
     end
   end

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

@@ -0,0 +1,135 @@
+# frozen_string_literal: true
+
+require "opentelemetry/sdk"
+require "json"
+require_relative "../../../../tokens"
+
+module Braintrust
+  module Trace
+    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
+            # @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 chat completions
+              wrap_chat(client, tracer_provider)
+
+              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]
+                        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::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
+                end
+              end
+
+              # Prepend the wrapper to the client's singleton class
+              client.singleton_class.prepend(wrapper)
+            end
+          end
+        end
+      end
+    end
+
+    # Backwards compatibility: this module was originally at Braintrust::Trace::AlexRudall::RubyOpenAI
+    module AlexRudall
+      RubyOpenAI = Contrib::Github::Alexrudall::RubyOpenAI
+    end
+  end
+end