braintrust 0.0.12 → 0.1.0

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

@@ -0,0 +1,298 @@
+# frozen_string_literal: true
+
+require "opentelemetry/sdk"
+require "json"
+
+require_relative "common"
+require_relative "../../../internal/time"
+require_relative "../../support/otel"
+require_relative "../../support/openai"
+
+module Braintrust
+  module Contrib
+    module OpenAI
+      module Instrumentation
+        # Chat completions instrumentation for OpenAI.
+        # Wraps create(), stream(), and stream_raw() methods to create spans.
+        module Chat
+          module Completions
+            def self.included(base)
+              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 create method for non-streaming completions
+              def create(**params)
+                client = instance_variable_get(:@client)
+                tracer = Braintrust::Contrib.tracer_for(client)
+
+                tracer.in_span("Chat Completion") do |span|
+                  metadata = build_metadata(params)
+
+                  set_input(span, params)
+
+                  response = nil
+                  time_to_first_token = Braintrust::Internal::Time.measure do
+                    response = super
+                  end
+
+                  set_output(span, response)
+                  set_metrics(span, response, time_to_first_token)
+                  finalize_metadata(span, metadata, response)
+
+                  response
+                end
+              end
+
+              # Wrap stream_raw for streaming chat completions (returns Internal::Stream)
+              # Stores context on stream object for span creation during consumption
+              def stream_raw(**params)
+                client = instance_variable_get(:@client)
+                tracer = Braintrust::Contrib.tracer_for(client)
+                metadata = build_metadata(params, stream: true)
+
+                stream_obj = super
+                Braintrust::Contrib::Context.set!(stream_obj,
+                  tracer: tracer,
+                  params: params,
+                  metadata: metadata,
+                  completions_instance: self,
+                  stream_type: :raw)
+                stream_obj
+              end
+
+              # Wrap stream for streaming chat completions (returns ChatCompletionStream)
+              # Stores context on stream object for span creation during consumption
+              def stream(**params)
+                client = instance_variable_get(:@client)
+                tracer = Braintrust::Contrib.tracer_for(client)
+                metadata = build_metadata(params, stream: true)
+
+                stream_obj = super
+                Braintrust::Contrib::Context.set!(stream_obj,
+                  tracer: tracer,
+                  params: params,
+                  metadata: metadata,
+                  completions_instance: self,
+                  stream_type: :chat_completion)
+                stream_obj
+              end
+
+              private
+
+              def build_metadata(params, stream: false)
+                metadata = {
+                  "provider" => "openai",
+                  "endpoint" => "/v1/chat/completions"
+                }
+                metadata["stream"] = true if stream
+                Completions::METADATA_FIELDS.each do |field|
+                  metadata[field.to_s] = params[field] if params.key?(field)
+                end
+                metadata
+              end
+
+              def set_input(span, params)
+                return unless params[:messages]
+
+                messages_array = params[:messages].map(&:to_h)
+                Support::OTel.set_json_attr(span, "braintrust.input_json", messages_array)
+              end
+
+              def set_output(span, response)
+                return unless response.respond_to?(:choices) && response.choices&.any?
+
+                choices_array = response.choices.map(&:to_h)
+                Support::OTel.set_json_attr(span, "braintrust.output_json", choices_array)
+              end
+
+              def set_metrics(span, response, time_to_first_token)
+                metrics = {}
+                if response.respond_to?(:usage) && response.usage
+                  metrics = Support::OpenAI.parse_usage_tokens(response.usage)
+                end
+                metrics["time_to_first_token"] = time_to_first_token
+                Support::OTel.set_json_attr(span, "braintrust.metrics", metrics) unless metrics.empty?
+              end
+
+              def finalize_metadata(span, metadata, response)
+                metadata["id"] = response.id if response.respond_to?(:id) && response.id
+                metadata["created"] = response.created if response.respond_to?(:created) && response.created
+                metadata["model"] = response.model if response.respond_to?(:model) && response.model
+                metadata["system_fingerprint"] = response.system_fingerprint if response.respond_to?(:system_fingerprint) && response.system_fingerprint
+                metadata["service_tier"] = response.service_tier if response.respond_to?(:service_tier) && response.service_tier
+                Support::OTel.set_json_attr(span, "braintrust.metadata", metadata)
+              end
+            end
+          end
+
+          # Instrumentation for ChatCompletionStream (returned by stream())
+          # Uses current_completion_snapshot for accumulated output
+          module ChatCompletionStream
+            def self.included(base)
+              base.prepend(InstanceMethods) unless applied?(base)
+            end
+
+            def self.applied?(base)
+              base.ancestors.include?(InstanceMethods)
+            end
+
+            module InstanceMethods
+              def each(&block)
+                ctx = Braintrust::Contrib::Context.from(self)
+                return super unless ctx&.[](:tracer) && !ctx[:consumed]
+
+                trace_consumption(ctx) { super(&block) }
+              end
+
+              def text
+                ctx = Braintrust::Contrib::Context.from(self)
+                return super unless ctx&.[](:tracer) && !ctx[:consumed]
+
+                original_enum = super
+                Enumerator.new do |y|
+                  trace_consumption(ctx) do
+                    original_enum.each { |t| y << t }
+                  end
+                end
+              end
+
+              private
+
+              def trace_consumption(ctx)
+                ctx[:consumed] = true
+
+                tracer = ctx[:tracer]
+                params = ctx[:params]
+                metadata = ctx[:metadata]
+                completions_instance = ctx[:completions_instance]
+                start_time = Braintrust::Internal::Time.measure
+
+                tracer.in_span("Chat Completion") do |span|
+                  completions_instance.send(:set_input, span, params)
+                  Support::OTel.set_json_attr(span, "braintrust.metadata", metadata)
+
+                  yield
+
+                  finalize_stream_span(span, start_time, metadata, completions_instance)
+                end
+              end
+
+              def finalize_stream_span(span, start_time, metadata, completions_instance)
+                time_to_first_token = Braintrust::Internal::Time.measure(start_time)
+
+                begin
+                  snapshot = current_completion_snapshot
+                  return unless snapshot
+
+                  # Set output from accumulated choices
+                  if snapshot.choices&.any?
+                    choices_array = snapshot.choices.map(&:to_h)
+                    Support::OTel.set_json_attr(span, "braintrust.output_json", choices_array)
+                  end
+
+                  # Set metrics
+                  metrics = {}
+                  if snapshot.usage
+                    metrics = Support::OpenAI.parse_usage_tokens(snapshot.usage)
+                  end
+                  metrics["time_to_first_token"] = time_to_first_token
+                  Support::OTel.set_json_attr(span, "braintrust.metrics", metrics) unless metrics.empty?
+
+                  # Update metadata with response fields
+                  metadata["id"] = snapshot.id if snapshot.respond_to?(:id) && snapshot.id
+                  metadata["created"] = snapshot.created if snapshot.respond_to?(:created) && snapshot.created
+                  metadata["model"] = snapshot.model if snapshot.respond_to?(:model) && snapshot.model
+                  metadata["system_fingerprint"] = snapshot.system_fingerprint if snapshot.respond_to?(:system_fingerprint) && snapshot.system_fingerprint
+                  metadata["service_tier"] = snapshot.service_tier if snapshot.respond_to?(:service_tier) && snapshot.service_tier
+                  Support::OTel.set_json_attr(span, "braintrust.metadata", metadata)
+                rescue => e
+                  Braintrust::Log.debug("Failed to get completion snapshot: #{e.message}")
+                end
+              end
+            end
+          end
+
+          # Instrumentation for Internal::Stream (returned by stream_raw())
+          # Aggregates chunks manually since Internal::Stream has no built-in accumulation
+          module InternalStream
+            def self.included(base)
+              base.prepend(InstanceMethods) unless applied?(base)
+            end
+
+            def self.applied?(base)
+              base.ancestors.include?(InstanceMethods)
+            end
+
+            module InstanceMethods
+              def each(&block)
+                ctx = Braintrust::Contrib::Context.from(self)
+                # Only trace if context present and is for chat completions (not other endpoints)
+                return super unless ctx&.[](:tracer) && !ctx[:consumed] && ctx[:stream_type] == :raw
+
+                ctx[:consumed] = true
+
+                tracer = ctx[:tracer]
+                params = ctx[:params]
+                metadata = ctx[:metadata]
+                completions_instance = ctx[:completions_instance]
+                aggregated_chunks = []
+                start_time = Braintrust::Internal::Time.measure
+                time_to_first_token = nil
+
+                tracer.in_span("Chat Completion") do |span|
+                  completions_instance.send(:set_input, span, params)
+                  Support::OTel.set_json_attr(span, "braintrust.metadata", metadata)
+
+                  super do |chunk|
+                    time_to_first_token ||= Braintrust::Internal::Time.measure(start_time)
+                    aggregated_chunks << chunk.to_h
+                    block&.call(chunk)
+                  end
+
+                  finalize_stream_span(span, aggregated_chunks, time_to_first_token, metadata)
+                end
+              end
+
+              private
+
+              def finalize_stream_span(span, aggregated_chunks, time_to_first_token, metadata)
+                return if aggregated_chunks.empty?
+
+                aggregated_output = Common.aggregate_streaming_chunks(aggregated_chunks)
+                Support::OTel.set_json_attr(span, "braintrust.output_json", aggregated_output[:choices])
+
+                # Set metrics
+                metrics = {}
+                if aggregated_output[:usage]
+                  metrics = Support::OpenAI.parse_usage_tokens(aggregated_output[:usage])
+                end
+                metrics["time_to_first_token"] = time_to_first_token
+                Support::OTel.set_json_attr(span, "braintrust.metrics", metrics) unless metrics.empty?
+
+                # Update metadata with response fields
+                metadata["id"] = aggregated_output[:id] if aggregated_output[:id]
+                metadata["created"] = aggregated_output[:created] if aggregated_output[:created]
+                metadata["model"] = aggregated_output[:model] if aggregated_output[:model]
+                metadata["system_fingerprint"] = aggregated_output[:system_fingerprint] if aggregated_output[:system_fingerprint]
+                metadata["service_tier"] = aggregated_output[:service_tier] if aggregated_output[:service_tier]
+                Support::OTel.set_json_attr(span, "braintrust.metadata", metadata)
+              end
+            end
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,134 @@
+# frozen_string_literal: true
+
+module Braintrust
+  module Contrib
+    module OpenAI
+      module Instrumentation
+        # Aggregation utilities for official OpenAI SDK instrumentation.
+        # These are specific to the official openai gem's data structures (symbol keys, SDK objects).
+        module Common
+          # Aggregate streaming chunks into a single response structure.
+          # Specific to official OpenAI SDK which uses symbol keys and SDK objects.
+          # @param chunks [Array<Hash>] array of chunk hashes from stream (symbol keys)
+          # @return [Hash] aggregated response with choices, usage, etc. (symbol keys)
+          def self.aggregate_streaming_chunks(chunks)
+            return {} if chunks.empty?
+
+            # Initialize aggregated structure
+            aggregated = {
+              id: nil,
+              created: nil,
+              model: nil,
+              system_fingerprint: nil,
+              choices: [],
+              usage: nil
+            }
+
+            # 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
+              next unless chunk[:choices].is_a?(Array)
+              chunk[: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
+                if delta[:tool_calls].is_a?(Array) && delta[:tool_calls].any?
+                  delta[:tool_calls].each do |tool_call_delta|
+                    if tool_call_delta[:id] && !tool_call_delta[:id].empty?
+                      # New tool call (dup strings to avoid mutating input)
+                      choice_data[index][:tool_calls] << {
+                        id: tool_call_delta[: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 events into a single response structure.
+          # Specific to official OpenAI SDK which returns typed event objects.
+          # @param events [Array] array of event objects from stream
+          # @return [Hash] aggregated response with output, usage, etc.
+          def self.aggregate_responses_events(events)
+            return {} if events.empty?
+
+            # Find the response.completed event which has the final response
+            completed_event = events.find { |e| e.respond_to?(:type) && e.type == :"response.completed" }
+
+            if completed_event&.respond_to?(:response)
+              response = completed_event.response
+              return {
+                id: response.respond_to?(:id) ? response.id : nil,
+                output: response.respond_to?(:output) ? response.output : nil,
+                usage: response.respond_to?(:usage) ? response.usage : nil
+              }
+            end
+
+            # Fallback if no completed event found
+            {}
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,187 @@
+# frozen_string_literal: true
+
+require "opentelemetry/sdk"
+require "json"
+
+require_relative "common"
+require_relative "../../../internal/time"
+require_relative "../../support/otel"
+require_relative "../../support/openai"
+
+module Braintrust
+  module Contrib
+    module OpenAI
+      module Instrumentation
+        # Responses API instrumentation for OpenAI.
+        # Wraps create() and stream() methods to create spans.
+        module Responses
+          def self.included(base)
+            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 non-streaming create method
+            def create(**params)
+              client = instance_variable_get(:@client)
+              tracer = Braintrust::Contrib.tracer_for(client)
+
+              tracer.in_span("openai.responses.create") do |span|
+                metadata = build_metadata(params)
+
+                set_input(span, params)
+
+                response = nil
+                time_to_first_token = Braintrust::Internal::Time.measure do
+                  response = super
+                end
+
+                set_output(span, response)
+                set_metrics(span, response, time_to_first_token)
+                finalize_metadata(span, metadata, response)
+
+                response
+              end
+            end
+
+            # Wrap streaming method
+            # Stores context on stream object for span creation during consumption
+            def stream(**params)
+              client = instance_variable_get(:@client)
+              tracer = Braintrust::Contrib.tracer_for(client)
+              metadata = build_metadata(params, stream: true)
+
+              stream_obj = super
+
+              Braintrust::Contrib::Context.set!(stream_obj,
+                tracer: tracer,
+                params: params,
+                metadata: metadata,
+                responses_instance: self)
+              stream_obj
+            end
+
+            private
+
+            def build_metadata(params, stream: false)
+              metadata = {
+                "provider" => "openai",
+                "endpoint" => "/v1/responses"
+              }
+              metadata["stream"] = true if stream
+              Responses::METADATA_FIELDS.each do |field|
+                metadata[field.to_s] = params[field] if params.key?(field)
+              end
+              metadata
+            end
+
+            def set_input(span, params)
+              return unless params[:input]
+
+              Support::OTel.set_json_attr(span, "braintrust.input_json", params[:input])
+            end
+
+            def set_output(span, response)
+              return unless response.respond_to?(:output) && response.output
+
+              Support::OTel.set_json_attr(span, "braintrust.output_json", response.output)
+            end
+
+            def set_metrics(span, response, time_to_first_token)
+              metrics = {}
+              if response.respond_to?(:usage) && response.usage
+                metrics = Support::OpenAI.parse_usage_tokens(response.usage)
+              end
+              metrics["time_to_first_token"] = time_to_first_token
+              Support::OTel.set_json_attr(span, "braintrust.metrics", metrics) unless metrics.empty?
+            end
+
+            def finalize_metadata(span, metadata, response)
+              metadata["id"] = response.id if response.respond_to?(:id) && response.id
+              Support::OTel.set_json_attr(span, "braintrust.metadata", metadata)
+            end
+          end
+        end
+
+        # Instrumentation for ResponseStream (returned by stream())
+        # Aggregates events and creates span lazily when consumed
+        module ResponseStream
+          def self.included(base)
+            base.prepend(InstanceMethods) unless applied?(base)
+          end
+
+          def self.applied?(base)
+            base.ancestors.include?(InstanceMethods)
+          end
+
+          module InstanceMethods
+            def each(&block)
+              ctx = Braintrust::Contrib::Context.from(self)
+              return super unless ctx&.[](:tracer) && !ctx[:consumed]
+
+              ctx[:consumed] = true
+
+              tracer = ctx[:tracer]
+              params = ctx[:params]
+              metadata = ctx[:metadata]
+              responses_instance = ctx[:responses_instance]
+              aggregated_events = []
+              start_time = Braintrust::Internal::Time.measure
+              time_to_first_token = nil
+
+              tracer.in_span("openai.responses.create") do |span|
+                responses_instance.send(:set_input, span, params)
+                Support::OTel.set_json_attr(span, "braintrust.metadata", metadata)
+
+                begin
+                  super do |event|
+                    time_to_first_token ||= Braintrust::Internal::Time.measure(start_time)
+                    aggregated_events << event
+                    block&.call(event)
+                  end
+                rescue => e
+                  span.record_exception(e)
+                  span.status = ::OpenTelemetry::Trace::Status.error("Streaming error: #{e.message}")
+                  raise
+                end
+
+                finalize_stream_span(span, aggregated_events, time_to_first_token, metadata)
+              end
+            end
+
+            private
+
+            def finalize_stream_span(span, aggregated_events, time_to_first_token, metadata)
+              return if aggregated_events.empty?
+
+              aggregated_output = Common.aggregate_responses_events(aggregated_events)
+              Support::OTel.set_json_attr(span, "braintrust.output_json", aggregated_output[:output]) if aggregated_output[:output]
+
+              # Set metrics
+              metrics = {}
+              if aggregated_output[:usage]
+                metrics = Support::OpenAI.parse_usage_tokens(aggregated_output[:usage])
+              end
+              metrics["time_to_first_token"] = time_to_first_token
+              Support::OTel.set_json_attr(span, "braintrust.metrics", metrics) unless metrics.empty?
+
+              # Update metadata with response fields
+              metadata["id"] = aggregated_output[:id] if aggregated_output[:id]
+              Support::OTel.set_json_attr(span, "braintrust.metadata", metadata)
+            end
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,58 @@
+# frozen_string_literal: true
+
+require_relative "../integration"
+require_relative "deprecated"
+
+module Braintrust
+  module Contrib
+    module OpenAI
+      # OpenAI integration for automatic instrumentation.
+      # Instruments the official openai gem (not ruby-openai).
+      class Integration
+        include Braintrust::Contrib::Integration
+
+        MINIMUM_VERSION = "0.1.0"
+
+        GEM_NAMES = ["openai"].freeze
+        REQUIRE_PATHS = ["openai"].freeze
+
+        # @return [Symbol] Unique identifier for this integration
+        def self.integration_name
+          :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 official openai gem is available
+        def self.loaded?
+          # Check if the official openai gem is loaded (not ruby-openai).
+          # The ruby-openai gem also uses "require 'openai'", so we need to distinguish them.
+
+          # This module is defined ONLY in the official OpenAI gem
+          defined?(::OpenAI::Internal) ? true : false
+        end
+
+        # Lazy-load the patcher only when actually patching.
+        # This keeps the integration stub lightweight.
+        # @return [Class] The patcher class
+        def self.patchers
+          require_relative "patcher"
+          [ChatPatcher, ResponsesPatcher]
+        end
+      end
+    end
+  end
+end