braintrust 0.0.12 → 0.1.1

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/moderations.rb

@@ -0,0 +1,93 @@
+# frozen_string_literal: true
+
+require "opentelemetry/sdk"
+require "json"
+
+require_relative "../../../internal/time"
+require_relative "../../support/otel"
+
+module Braintrust
+  module Contrib
+    module OpenAI
+      module Instrumentation
+        # Moderations API instrumentation for OpenAI.
+        # Wraps create() method to create spans.
+        module Moderations
+          def self.included(base)
+            base.prepend(InstanceMethods) unless applied?(base)
+          end
+
+          def self.applied?(base)
+            base.ancestors.include?(InstanceMethods)
+          end
+
+          METADATA_FIELDS = %i[
+            model
+          ].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.moderations.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, time_to_first_token)
+                finalize_metadata(span, metadata, response)
+
+                response
+              end
+            end
+
+            private
+
+            def build_metadata(params)
+              metadata = {
+                "provider" => "openai",
+                "endpoint" => "/v1/moderations"
+              }
+              Moderations::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?(:results) && response.results
+
+              Support::OTel.set_json_attr(span, "braintrust.output_json", response.results)
+            end
+
+            def set_metrics(span, time_to_first_token)
+              metrics = {}
+              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["model"] = response.model if response.respond_to?(:model) && response.model
+              Support::OTel.set_json_attr(span, "braintrust.metadata", metadata)
+            end
+          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, ModerationsPatcher]
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,173 @@
+# frozen_string_literal: true
+
+require_relative "../patcher"
+require_relative "instrumentation/chat"
+require_relative "instrumentation/responses"
+require_relative "instrumentation/moderations"
+
+module Braintrust
+  module Contrib
+    module OpenAI
+      # Patcher for OpenAI integration - implements class-level patching.
+      # All new OpenAI::Client instances created after patch! will be automatically instrumented.
+      class ChatPatcher < Braintrust::Contrib::Patcher
+        class << self
+          def applicable?
+            defined?(::OpenAI::Client)
+          end
+
+          def patched?(**options)
+            # Use the target's singleton class if provided, otherwise check the base class.
+            target_class = get_singleton_class(options[:target]) || ::OpenAI::Resources::Chat::Completions
+
+            Instrumentation::Chat::Completions.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?
+
+            # Stream classes are shared across all clients, patch at class level.
+            # The instrumentation short-circuits when no context is present,
+            # so uninstrumented clients' streams pass through unaffected.
+            patch_stream_classes
+
+            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)
+
+              get_singleton_class(options[:target]).include(Instrumentation::Chat::Completions)
+            else
+              # Class-level (for all clients)
+              ::OpenAI::Resources::Chat::Completions.include(Instrumentation::Chat::Completions)
+            end
+          end
+
+          def patch_stream_classes
+            # Patch ChatCompletionStream for stream() method
+            if defined?(::OpenAI::Helpers::Streaming::ChatCompletionStream)
+              unless Instrumentation::Chat::ChatCompletionStream.applied?(::OpenAI::Helpers::Streaming::ChatCompletionStream)
+                ::OpenAI::Helpers::Streaming::ChatCompletionStream.include(Instrumentation::Chat::ChatCompletionStream)
+              end
+            end
+
+            # Patch Internal::Stream for stream_raw() method
+            if defined?(::OpenAI::Internal::Stream)
+              unless Instrumentation::Chat::InternalStream.applied?(::OpenAI::Internal::Stream)
+                ::OpenAI::Internal::Stream.include(Instrumentation::Chat::InternalStream)
+              end
+            end
+          end
+
+          private
+
+          def get_singleton_class(client)
+            client&.chat&.completions&.singleton_class
+          end
+        end
+      end
+
+      # Patcher for OpenAI integration - implements class-level patching.
+      # All new OpenAI::Client instances created after patch! will be automatically instrumented.
+      class ResponsesPatcher < Braintrust::Contrib::Patcher
+        class << self
+          def applicable?
+            defined?(::OpenAI::Client) && ::OpenAI::Client.instance_methods.include?(:responses)
+          end
+
+          def patched?(**options)
+            # Use the target's singleton class if provided, otherwise check the base class.
+            target_class = get_singleton_class(options[:target]) || ::OpenAI::Resources::Responses
+
+            Instrumentation::Responses.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?
+
+            # Stream class is shared across all clients, patch at class level.
+            # The instrumentation short-circuits when no context is present,
+            # so uninstrumented clients' streams pass through unaffected.
+            patch_response_stream
+
+            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)
+
+              get_singleton_class(options[:target]).include(Instrumentation::Responses)
+            else
+              # Class-level (for all clients)
+              ::OpenAI::Resources::Responses.include(Instrumentation::Responses)
+            end
+          end
+
+          def patch_response_stream
+            # Patch ResponseStream for stream() method
+            if defined?(::OpenAI::Helpers::Streaming::ResponseStream)
+              unless Instrumentation::ResponseStream.applied?(::OpenAI::Helpers::Streaming::ResponseStream)
+                ::OpenAI::Helpers::Streaming::ResponseStream.include(Instrumentation::ResponseStream)
+              end
+            end
+          end
+
+          private
+
+          def get_singleton_class(client)
+            client&.responses&.singleton_class
+          end
+        end
+      end
+
+      # Patcher for OpenAI Moderations API - implements class-level patching.
+      # All new OpenAI::Client instances created after patch! will be automatically instrumented.
+      class ModerationsPatcher < Braintrust::Contrib::Patcher
+        class << self
+          def applicable?
+            defined?(::OpenAI::Client) && ::OpenAI::Client.instance_methods.include?(:moderations)
+          end
+
+          def patched?(**options)
+            # Use the target's singleton class if provided, otherwise check the base class.
+            target_class = get_singleton_class(options[:target]) || ::OpenAI::Resources::Moderations
+
+            Instrumentation::Moderations.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)
+
+              get_singleton_class(options[:target]).include(Instrumentation::Moderations)
+            else
+              # Class-level (for all clients)
+              ::OpenAI::Resources::Moderations.include(Instrumentation::Moderations)
+            end
+          end
+
+          private
+
+          def get_singleton_class(client)
+            client&.moderations&.singleton_class
+          end
+        end
+      end
+    end
+  end
+end