braintrust 0.0.12 → 0.1.0

data/lib/braintrust/contrib/ruby_openai/instrumentation/chat.rb

@@ -0,0 +1,149 @@
+# frozen_string_literal: true
+
+require "opentelemetry/sdk"
+require "json"
+
+require_relative "common"
+require_relative "../../support/otel"
+require_relative "../../support/openai"
+require_relative "../../../internal/time"
+
+module Braintrust
+  module Contrib
+    module RubyOpenAI
+      module Instrumentation
+        # Chat completions instrumentation for ruby-openai.
+        # Provides module that can be prepended to OpenAI::Client to instrument the chat method.
+        module Chat
+          def self.included(base)
+            # Guard against double-wrapping: Check if patch is already in the ancestor chain.
+            # This prevents double instrumentation if class-level patching was already applied,
+            # and this patch is being applied to a singleton-class. (Special case.)
+            #
+            # Ruby's prepend() doesn't check the full inheritance chain, so without this guard,
+            # the instrumentation could be added twice.
+            base.prepend(InstanceMethods) unless applied?(base)
+          end
+
+          def self.applied?(base)
+            base.ancestors.include?(InstanceMethods)
+          end
+
+          METADATA_FIELDS = %i[
+            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
+          ].freeze
+
+          module InstanceMethods
+            # Wrap chat method for ruby-openai gem
+            # ruby-openai API: client.chat(parameters: {...})
+            def chat(parameters:)
+              tracer = Braintrust::Contrib.tracer_for(self)
+
+              tracer.in_span("Chat Completion") do |span|
+                is_streaming = streaming?(parameters)
+                metadata = build_metadata(parameters)
+                set_input(span, parameters)
+
+                aggregated_chunks = []
+                time_to_first_token = nil
+                response = nil
+                response_data = {}
+
+                if is_streaming
+                  # Setup a time measurement for the first chunk from the stream
+                  start_time = nil
+                  parameters = wrap_stream_callback(parameters, aggregated_chunks) do
+                    time_to_first_token ||= Braintrust::Internal::Time.measure(start_time)
+                  end
+                  start_time = Braintrust::Internal::Time.measure
+
+                  # Then initiate the stream
+                  response = super(parameters: parameters)
+
+                  if !aggregated_chunks.empty?
+                    response_data = Common.aggregate_streaming_chunks(aggregated_chunks)
+                  end
+                else
+                  # Make a time measurement synchronously around the API call
+                  time_to_first_token = Braintrust::Internal::Time.measure do
+                    response = super(parameters: parameters)
+                    response_data = response if response
+                  end
+                end
+
+                set_output(span, response_data)
+                set_metrics(span, response_data, time_to_first_token)
+                finalize_metadata(span, metadata, response_data)
+
+                response
+              end
+            end
+
+            private
+
+            def streaming?(parameters)
+              parameters.key?(:stream) && parameters[:stream].is_a?(Proc)
+            end
+
+            def wrap_stream_callback(parameters, aggregated_chunks)
+              original_stream_proc = parameters[:stream]
+              parameters = parameters.dup
+
+              parameters[:stream] = proc do |chunk, bytesize|
+                yield if aggregated_chunks.empty?
+                aggregated_chunks << chunk
+                original_stream_proc.call(chunk, bytesize)
+              end
+
+              parameters
+            end
+
+            def build_metadata(parameters)
+              metadata = {
+                "provider" => "openai",
+                "endpoint" => "/v1/chat/completions"
+              }
+
+              Chat::METADATA_FIELDS.each do |field|
+                next unless parameters.key?(field)
+                # Stream param is a Proc - just mark as true
+                metadata[field.to_s] = (field == :stream) ? true : parameters[field]
+              end
+
+              metadata
+            end
+
+            def set_input(span, parameters)
+              return unless parameters[:messages]
+              Support::OTel.set_json_attr(span, "braintrust.input_json", parameters[:messages])
+            end
+
+            def set_output(span, response_data)
+              choices = response_data[:choices] || response_data["choices"]
+              return unless choices&.any?
+              Support::OTel.set_json_attr(span, "braintrust.output_json", choices)
+            end
+
+            def set_metrics(span, response_data, time_to_first_token)
+              usage = response_data[:usage] || response_data["usage"]
+              metrics = usage ? Support::OpenAI.parse_usage_tokens(usage) : {}
+              metrics["time_to_first_token"] = time_to_first_token || 0.0
+              Support::OTel.set_json_attr(span, "braintrust.metrics", metrics) unless metrics.empty?
+            end
+
+            def finalize_metadata(span, metadata, response_data)
+              %w[id created model system_fingerprint service_tier].each do |field|
+                value = response_data[field.to_sym] || response_data[field]
+                metadata[field] = value if value
+              end
+              Support::OTel.set_json_attr(span, "braintrust.metadata", metadata)
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/braintrust/contrib/ruby_openai/instrumentation/common.rb

@@ -0,0 +1,138 @@
+# frozen_string_literal: true
+
+module Braintrust
+  module Contrib
+    module RubyOpenAI
+      module Instrumentation
+        # Aggregation utilities for ruby-openai gem instrumentation.
+        # These are specific to the ruby-openai gem's data structures (string keys, plain hashes).
+        module Common
+          # Aggregate streaming chunks into a single response structure.
+          # Specific to ruby-openai gem which uses string keys and plain hashes.
+          # @param chunks [Array<Hash>] array of chunk hashes from stream (string keys)
+          # @return [Hash] aggregated response with choices, usage, etc. (string keys)
+          def self.aggregate_streaming_chunks(chunks)
+            return {} if chunks.empty?
+
+            # Initialize aggregated structure
+            aggregated = {
+              "id" => nil,
+              "created" => nil,
+              "model" => nil,
+              "system_fingerprint" => nil,
+              "usage" => nil,
+              "choices" => []
+            }
+
+            # Track aggregated content and tool_calls for each choice index
+            choice_data = {}
+
+            chunks.each do |chunk|
+              # Capture top-level fields from any chunk that has them
+              aggregated["id"] ||= chunk["id"]
+              aggregated["created"] ||= chunk["created"]
+              aggregated["model"] ||= chunk["model"]
+              aggregated["system_fingerprint"] ||= chunk["system_fingerprint"]
+
+              # Aggregate usage (usually only in last chunk if stream_options.include_usage is set)
+              aggregated["usage"] = chunk["usage"] if chunk["usage"]
+
+              # Process choices
+              choices = chunk["choices"]
+              next unless choices.is_a?(Array)
+
+              choices.each do |choice|
+                index = choice["index"] || 0
+                choice_data[index] ||= {
+                  "index" => index,
+                  "role" => nil,
+                  "content" => +"",
+                  "tool_calls" => [],
+                  "finish_reason" => nil
+                }
+
+                delta = choice["delta"] || {}
+
+                # Aggregate role (set once from first delta that has it)
+                choice_data[index]["role"] ||= delta["role"]
+
+                # Aggregate content
+                choice_data[index]["content"] << delta["content"] if delta["content"]
+
+                # Aggregate tool_calls
+                tool_calls = delta["tool_calls"]
+                if tool_calls.is_a?(Array) && tool_calls.any?
+                  tool_calls.each do |tool_call_delta|
+                    tc_id = tool_call_delta["id"]
+                    if tc_id && !tc_id.empty?
+                      # New tool call
+                      choice_data[index]["tool_calls"] << {
+                        "id" => tc_id,
+                        "type" => tool_call_delta["type"],
+                        "function" => {
+                          "name" => +(tool_call_delta.dig("function", "name") || ""),
+                          "arguments" => +(tool_call_delta.dig("function", "arguments") || "")
+                        }
+                      }
+                    elsif choice_data[index]["tool_calls"].any?
+                      # Continuation - append arguments to last tool call
+                      last_tool_call = choice_data[index]["tool_calls"].last
+                      if tool_call_delta.dig("function", "arguments")
+                        last_tool_call["function"]["arguments"] << tool_call_delta["function"]["arguments"]
+                      end
+                    end
+                  end
+                end
+
+                # Capture finish_reason
+                choice_data[index]["finish_reason"] = choice["finish_reason"] if choice["finish_reason"]
+              end
+            end
+
+            # Build final choices array
+            aggregated["choices"] = choice_data.values.sort_by { |c| c["index"] }.map do |choice|
+              message = {
+                "role" => choice["role"],
+                "content" => choice["content"].empty? ? nil : choice["content"]
+              }
+
+              # Add tool_calls to message if any
+              message["tool_calls"] = choice["tool_calls"] if choice["tool_calls"].any?
+
+              {
+                "index" => choice["index"],
+                "message" => message,
+                "finish_reason" => choice["finish_reason"]
+              }
+            end
+
+            aggregated
+          end
+
+          # Aggregate responses streaming chunks into a single response structure.
+          # Specific to ruby-openai gem which uses string keys and plain hashes.
+          # @param chunks [Array<Hash>] array of chunk hashes from stream (string keys)
+          # @return [Hash] aggregated response with output, usage, id (string keys)
+          def self.aggregate_responses_chunks(chunks)
+            return {} if chunks.empty?
+
+            # Find the response.completed event which has the final response
+            completed_chunk = chunks.find { |c| c["type"] == "response.completed" }
+
+            if completed_chunk && completed_chunk["response"]
+              response = completed_chunk["response"]
+              return {
+                "id" => response["id"],
+                "output" => response["output"],
+                "usage" => response["usage"]
+              }
+            end
+
+            # Fallback if no completed event found
+            {}
+          end
+        end
+      end
+    end
+  end
+end

data/lib/braintrust/contrib/ruby_openai/instrumentation/responses.rb

@@ -0,0 +1,146 @@
+# frozen_string_literal: true
+
+require "opentelemetry/sdk"
+require "json"
+
+require_relative "common"
+require_relative "../../support/otel"
+require_relative "../../support/openai"
+require_relative "../../../internal/time"
+
+module Braintrust
+  module Contrib
+    module RubyOpenAI
+      module Instrumentation
+        # Responses API instrumentation for ruby-openai.
+        # Provides module that can be prepended to OpenAI::Responses to instrument the create method.
+        module Responses
+          def self.included(base)
+            # Guard against double-wrapping: Check if patch is already in the ancestor chain.
+            # This prevents double instrumentation if class-level patching was already applied,
+            # and this patch is being applied to a singleton-class. (Special case.)
+            #
+            # Ruby's prepend() doesn't check the full inheritance chain, so without this guard,
+            # the instrumentation could be added twice.
+            base.prepend(InstanceMethods) unless applied?(base)
+          end
+
+          def self.applied?(base)
+            base.ancestors.include?(InstanceMethods)
+          end
+
+          METADATA_FIELDS = %i[
+            model instructions modalities tools parallel_tool_calls
+            tool_choice temperature max_tokens top_p frequency_penalty
+            presence_penalty seed user metadata store response_format
+            reasoning previous_response_id truncation
+          ].freeze
+
+          module InstanceMethods
+            # Wrap create method for ruby-openai responses API
+            # ruby-openai API: client.responses.create(parameters: {...})
+            def create(parameters:)
+              client = instance_variable_get(:@client)
+              tracer = Braintrust::Contrib.tracer_for(client)
+
+              tracer.in_span("openai.responses.create") do |span|
+                is_streaming = streaming?(parameters)
+                metadata = build_metadata(parameters)
+                set_input(span, parameters)
+
+                aggregated_chunks = []
+                time_to_first_token = nil
+                response = nil
+                response_data = {}
+
+                if is_streaming
+                  # Setup a time measurement for the first chunk from the stream
+                  start_time = nil
+                  parameters = wrap_stream_callback(parameters, aggregated_chunks) do
+                    time_to_first_token ||= Braintrust::Internal::Time.measure(start_time)
+                  end
+                  start_time = Braintrust::Internal::Time.measure
+
+                  # Then initiate the stream
+                  response = super(parameters: parameters)
+
+                  if !aggregated_chunks.empty?
+                    response_data = Common.aggregate_responses_chunks(aggregated_chunks)
+                  end
+                else
+                  # Make a time measurement synchronously around the API call
+                  time_to_first_token = Braintrust::Internal::Time.measure do
+                    response = super(parameters: parameters)
+                    response_data = response if response
+                  end
+                end
+
+                set_output(span, response_data)
+                set_metrics(span, response_data, time_to_first_token)
+                finalize_metadata(span, metadata, response_data)
+
+                response
+              end
+            end
+
+            private
+
+            def streaming?(parameters)
+              parameters.key?(:stream) && parameters[:stream].is_a?(Proc)
+            end
+
+            def wrap_stream_callback(parameters, aggregated_chunks)
+              original_stream_proc = parameters[:stream]
+              parameters = parameters.dup
+
+              parameters[:stream] = proc do |chunk, event|
+                yield if aggregated_chunks.empty?
+                aggregated_chunks << chunk
+                original_stream_proc.call(chunk, event)
+              end
+
+              parameters
+            end
+
+            def build_metadata(parameters)
+              metadata = {
+                "provider" => "openai",
+                "endpoint" => "/v1/responses"
+              }
+
+              Responses::METADATA_FIELDS.each do |field|
+                metadata[field.to_s] = parameters[field] if parameters.key?(field)
+              end
+
+              metadata["stream"] = true if streaming?(parameters)
+              metadata
+            end
+
+            def set_input(span, parameters)
+              return unless parameters[:input]
+              Support::OTel.set_json_attr(span, "braintrust.input_json", parameters[:input])
+            end
+
+            def set_output(span, response_data)
+              output = response_data["output"]
+              return unless output
+              Support::OTel.set_json_attr(span, "braintrust.output_json", output)
+            end
+
+            def set_metrics(span, response_data, time_to_first_token)
+              usage = response_data["usage"]
+              metrics = usage ? Support::OpenAI.parse_usage_tokens(usage) : {}
+              metrics["time_to_first_token"] = time_to_first_token || 0.0
+              Support::OTel.set_json_attr(span, "braintrust.metrics", metrics) unless metrics.empty?
+            end
+
+            def finalize_metadata(span, metadata, response_data)
+              metadata["id"] = response_data["id"] if response_data["id"]
+              Support::OTel.set_json_attr(span, "braintrust.metadata", metadata)
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/braintrust/contrib/ruby_openai/integration.rb

@@ -0,0 +1,58 @@
+# frozen_string_literal: true
+
+require_relative "../integration"
+require_relative "deprecated"
+
+module Braintrust
+  module Contrib
+    module RubyOpenAI
+      # RubyOpenAI integration for automatic instrumentation.
+      # Instruments the alexrudall ruby-openai gem (not the official openai gem).
+      class Integration
+        include Braintrust::Contrib::Integration
+
+        MINIMUM_VERSION = "7.0.0"
+
+        GEM_NAMES = ["ruby-openai"].freeze
+        REQUIRE_PATHS = ["openai"].freeze
+
+        # @return [Symbol] Unique identifier for this integration
+        def self.integration_name
+          :ruby_openai
+        end
+
+        # @return [Array<String>] Gem names this integration supports
+        def self.gem_names
+          GEM_NAMES
+        end
+
+        # @return [Array<String>] Require paths for auto-instrument detection
+        def self.require_paths
+          REQUIRE_PATHS
+        end
+
+        # @return [String] Minimum compatible version
+        def self.minimum_version
+          MINIMUM_VERSION
+        end
+
+        # @return [Boolean] true if ruby-openai gem is available (not official openai gem)
+        def self.loaded?
+          # Check if ruby-openai gem is loaded (not the official openai gem).
+          # Both gems use "require 'openai'", so we need to distinguish them.
+          #
+          # OpenAI::Internal is defined ONLY in the official OpenAI gem
+          (defined?(::OpenAI::Client) && !defined?(::OpenAI::Internal)) ? true : false
+        end
+
+        # Lazy-load the patcher only when actually patching.
+        # This keeps the integration stub lightweight.
+        # @return [Array<Class>] The patcher classes
+        def self.patchers
+          require_relative "patcher"
+          [ChatPatcher, ResponsesPatcher]
+        end
+      end
+    end
+  end
+end

data/lib/braintrust/contrib/ruby_openai/patcher.rb

@@ -0,0 +1,85 @@
+# frozen_string_literal: true
+
+require_relative "../patcher"
+require_relative "instrumentation/chat"
+require_relative "instrumentation/responses"
+
+module Braintrust
+  module Contrib
+    module RubyOpenAI
+      # Patcher for ruby-openai chat completions.
+      # Instruments OpenAI::Client#chat method.
+      class ChatPatcher < Braintrust::Contrib::Patcher
+        class << self
+          def applicable?
+            defined?(::OpenAI::Client)
+          end
+
+          def patched?(**options)
+            target_class = options[:target]&.singleton_class || ::OpenAI::Client
+            Instrumentation::Chat.applied?(target_class)
+          end
+
+          # Perform the actual patching.
+          # @param options [Hash] Configuration options passed from integration
+          # @option options [Object] :target Optional target instance to patch
+          # @option options [OpenTelemetry::SDK::Trace::TracerProvider] :tracer_provider Optional tracer provider
+          # @return [void]
+          def perform_patch(**options)
+            return unless applicable?
+
+            if options[:target]
+              # Instance-level (for only this client)
+              raise ArgumentError, "target must be a kind of ::OpenAI::Client" unless options[:target].is_a?(::OpenAI::Client)
+
+              options[:target].singleton_class.include(Instrumentation::Chat)
+            else
+              # Class-level (for all clients)
+              ::OpenAI::Client.include(Instrumentation::Chat)
+            end
+          end
+        end
+      end
+
+      # Patcher for ruby-openai responses API.
+      # Instruments OpenAI::Responses#create method.
+      class ResponsesPatcher < Braintrust::Contrib::Patcher
+        class << self
+          def applicable?
+            defined?(::OpenAI::Client) && ::OpenAI::Client.method_defined?(:responses)
+          end
+
+          def patched?(**options)
+            if options[:target]
+              responses_obj = options[:target].responses
+              Instrumentation::Responses.applied?(responses_obj.singleton_class)
+            else
+              # For class-level, check if the responses class is patched
+              defined?(::OpenAI::Responses) && Instrumentation::Responses.applied?(::OpenAI::Responses)
+            end
+          end
+
+          # Perform the actual patching.
+          # @param options [Hash] Configuration options passed from integration
+          # @option options [Object] :target Optional target instance to patch
+          # @option options [OpenTelemetry::SDK::Trace::TracerProvider] :tracer_provider Optional tracer provider
+          # @return [void]
+          def perform_patch(**options)
+            return unless applicable?
+
+            if options[:target]
+              # Instance-level (for only this client)
+              raise ArgumentError, "target must be a kind of ::OpenAI::Client" unless options[:target].is_a?(::OpenAI::Client)
+
+              responses_obj = options[:target].responses
+              responses_obj.singleton_class.include(Instrumentation::Responses)
+            else
+              # Class-level (for all clients)
+              ::OpenAI::Responses.include(Instrumentation::Responses)
+            end
+          end
+        end
+      end
+    end
+  end
+end