braintrust 0.0.12 → 0.1.1

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

@@ -0,0 +1,145 @@
+# frozen_string_literal: true
+
+require_relative "../patcher"
+require_relative "instrumentation/messages"
+require_relative "instrumentation/beta_messages"
+
+module Braintrust
+  module Contrib
+    module Anthropic
+      # Patcher for Anthropic messages.
+      # Instruments Anthropic::Messages#create and #stream methods.
+      class MessagesPatcher < Braintrust::Contrib::Patcher
+        class << self
+          def applicable?
+            defined?(::Anthropic::Client)
+          end
+
+          def patched?(**options)
+            target_class = get_singleton_class(options[:target]) || ::Anthropic::Resources::Messages
+            Instrumentation::Messages.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?
+
+            # MessageStream is shared across all clients, so patch at class level.
+            # The instrumentation short-circuits when no context is present,
+            # so uninstrumented clients' streams pass through unaffected.
+            patch_message_stream
+
+            if options[:target]
+              # Instance-level (for only this client instance)
+              raise ArgumentError, "target must be a kind of ::Anthropic::Client" unless options[:target].is_a?(::Anthropic::Client)
+
+              get_singleton_class(options[:target]).include(Instrumentation::Messages)
+            else
+              # Class-level (for all client instances)
+              ::Anthropic::Resources::Messages.include(Instrumentation::Messages)
+            end
+          end
+
+          private
+
+          def get_singleton_class(client)
+            client&.messages&.singleton_class
+          end
+
+          def patch_message_stream
+            return unless defined?(::Anthropic::Helpers::Streaming::MessageStream)
+            return if Instrumentation::MessageStream.applied?(::Anthropic::Helpers::Streaming::MessageStream)
+
+            ::Anthropic::Helpers::Streaming::MessageStream.include(Instrumentation::MessageStream)
+          end
+        end
+      end
+
+      # Patcher for Anthropic beta messages API.
+      # Instruments client.beta.messages.create and stream methods.
+      #
+      # @note Beta APIs are experimental and subject to change between SDK versions.
+      #   Braintrust will make reasonable efforts to maintain compatibility, but
+      #   breaking changes may require SDK updates.
+      #
+      # @see https://docs.anthropic.com/en/docs/build-with-claude/structured-outputs
+      #   for structured outputs documentation
+      class BetaMessagesPatcher < Braintrust::Contrib::Patcher
+        # Version constraints for beta patcher.
+        # Set MAXIMUM_VERSION when a breaking change is discovered to disable
+        # beta instrumentation on incompatible versions until a fix is released.
+        # Currently nil = rely on class existence check only.
+        MAXIMUM_VERSION = nil
+
+        class << self
+          def applicable?
+            return false unless defined?(::Anthropic::Client)
+            return false unless defined?(::Anthropic::Resources::Beta::Messages)
+            return false if MAXIMUM_VERSION && !version_compatible?
+            true
+          end
+
+          def patched?(**options)
+            target_class = get_singleton_class(options[:target]) || ::Anthropic::Resources::Beta::Messages
+            Instrumentation::BetaMessages.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?
+
+            Braintrust::Log.debug("Instrumenting Anthropic beta.messages API (experimental)")
+
+            # MessageStream is shared with stable API - already patched by MessagesPatcher.
+            # The BetaMessages instrumentation sets api_version: "beta" in context,
+            # which MessageStream uses to include in metadata.
+            patch_message_stream
+
+            if options[:target]
+              # Instance-level (for only this client instance)
+              raise ArgumentError, "target must be a kind of ::Anthropic::Client" unless options[:target].is_a?(::Anthropic::Client)
+
+              get_singleton_class(options[:target]).include(Instrumentation::BetaMessages)
+            else
+              # Class-level (for all client instances)
+              ::Anthropic::Resources::Beta::Messages.include(Instrumentation::BetaMessages)
+            end
+          end
+
+          private
+
+          def version_compatible?
+            return true unless MAXIMUM_VERSION
+
+            spec = Gem.loaded_specs["anthropic"]
+            return true unless spec
+
+            spec.version <= Gem::Version.new(MAXIMUM_VERSION)
+          end
+
+          def get_singleton_class(client)
+            client&.beta&.messages&.singleton_class
+          rescue
+            # Defensive: beta namespace may not exist or may have changed
+            nil
+          end
+
+          def patch_message_stream
+            return unless defined?(::Anthropic::Helpers::Streaming::MessageStream)
+            return if Instrumentation::MessageStream.applied?(::Anthropic::Helpers::Streaming::MessageStream)
+
+            ::Anthropic::Helpers::Streaming::MessageStream.include(Instrumentation::MessageStream)
+          end
+        end
+      end
+    end
+  end
+end

data/lib/braintrust/contrib/context.rb

@@ -0,0 +1,56 @@
+# frozen_string_literal: true
+
+module Braintrust
+  module Contrib
+    # Per-instance or per-class configuration context.
+    # Allows attaching generic configuration to specific objects or classes.
+    class Context
+      # Set or update context on a target object.
+      # Creates a new context if one doesn't exist, or updates existing context.
+      # @param target [Object] The object to attach context to
+      # @param options [Hash] Configuration options to store
+      # @return [Context, nil] The existing context if updated, nil if created new or options empty
+      def self.set!(target, **options)
+        return nil if options.empty?
+
+        if (ctx = from(target))
+          # Update existing context
+          options.each { |k, v| ctx[k] = v }
+        else
+          # Create and attach new context
+          target.instance_variable_set(:@braintrust_context, new(**options))
+        end
+
+        ctx
+      end
+
+      # Retrieve context from a target.
+      # @param target [Object] The object to retrieve context from
+      # @return [Context, nil] The context if found, nil otherwise
+      def self.from(target)
+        target&.instance_variable_get(:@braintrust_context)
+      end
+
+      # @param options [Hash] Configuration options
+      def initialize(**options)
+        @options = options
+      end
+
+      def [](key)
+        @options[key]
+      end
+
+      def []=(key, value)
+        @options[key] = value
+      end
+
+      # Get an option value with a default fallback.
+      # @param key [Symbol, String] The option key
+      # @param default [Object] The default value if key not found
+      # @return [Object] The option value, or default if not found
+      def fetch(key, default)
+        @options.fetch(key, default)
+      end
+    end
+  end
+end

data/lib/braintrust/contrib/integration.rb

@@ -0,0 +1,160 @@
+# frozen_string_literal: true
+
+module Braintrust
+  module Contrib
+    # Base module defining the integration contract.
+    # Include this module in integration classes to define the schema.
+    # Delegates actual patching to a Patcher subclass.
+    module Integration
+      def self.included(base)
+        base.extend(ClassMethods)
+      end
+
+      module ClassMethods
+        # Unique symbol name for this integration (e.g., :openai, :anthropic).
+        # @return [Symbol]
+        def integration_name
+          raise NotImplementedError, "#{self} must implement integration_name"
+        end
+
+        # Array of gem names this integration supports.
+        # @return [Array<String>]
+        def gem_names
+          raise NotImplementedError, "#{self} must implement gem_names"
+        end
+
+        # Require paths for auto-instrument detection.
+        # Default implementation returns gem_names.
+        # @return [Array<String>]
+        def require_paths
+          gem_names
+        end
+
+        # Is the target library available for loading?
+        # @return [Boolean]
+        def available?
+          gem_names.any? { |name| Gem.loaded_specs.key?(name) }
+        end
+
+        # Is the target library loaded?
+        # @return [Boolean]
+        def loaded?
+          raise NotImplementedError, "#{self} must implement loaded?"
+        end
+
+        # Minimum compatible version (optional, inclusive).
+        # @return [String, nil]
+        def minimum_version
+          nil
+        end
+
+        # Maximum compatible version (optional, inclusive).
+        # @return [String, nil]
+        def maximum_version
+          nil
+        end
+
+        # Is the library version compatible?
+        # @return [Boolean]
+        def compatible?
+          return false unless available?
+
+          gem_names.each do |name|
+            spec = Gem.loaded_specs[name]
+            next unless spec
+
+            version = spec.version
+            return false if minimum_version && version < Gem::Version.new(minimum_version)
+            return false if maximum_version && version > Gem::Version.new(maximum_version)
+            return true
+          end
+          false
+        end
+
+        # Array of patcher classes for this integration.
+        # Override to return multiple patchers for version-specific logic.
+        # @return [Array<Class>] Array of patcher classes
+        def patchers
+          [patcher] # Default: single patcher
+        end
+
+        # Convenience method for single patcher (existing pattern).
+        # Override this OR patchers (not both).
+        # @return [Class] The patcher class
+        def patcher
+          raise NotImplementedError, "#{self} must implement patcher or patchers"
+        end
+
+        # Instrument this integration with optional configuration.
+        # If a target is provided, configures the target instance specifically.
+        # Otherwise, applies class-level instrumentation to all instances.
+        #
+        # @param options [Hash] Configuration options
+        # @option options [Object] :target Optional target instance to instrument
+        # @option options [OpenTelemetry::SDK::Trace::TracerProvider] :tracer_provider Optional tracer provider
+        # @return [Boolean] true if patching succeeded or was already done
+        #
+        # @example Class-level instrumentation (all clients)
+        #   integration.instrument!(tracer_provider: my_provider)
+        #
+        # @example Instance-level instrumentation (specific client)
+        #   integration.instrument!(target: client, tracer_provider: my_provider)
+        def instrument!(**options)
+          if options.empty?
+            Braintrust::Log.debug("#{integration_name}.instrument! called")
+          else
+            Braintrust::Log.debug("#{integration_name}.instrument! called (#{options.keys.join(", ")})")
+          end
+
+          if options[:target]
+            # Configure the target with provided options (exclude :target from context)
+            context_options = options.except(:target)
+            Contrib::Context.set!(options[:target], **context_options) unless context_options.empty?
+          end
+
+          patch!(**options)
+        end
+
+        # Apply instrumentation (idempotent). Tries all applicable patchers.
+        # This method is typically called by instrument! after configuration.
+        #
+        # @param options [Hash] Configuration options
+        # @option options [Object] :target Optional target instance to patch
+        # @option options [OpenTelemetry::SDK::Trace::TracerProvider] :tracer_provider Optional tracer provider
+        # @return [Boolean] true if any patching succeeded or was already done
+        def patch!(**options)
+          unless available?
+            Braintrust::Log.debug("#{integration_name}.patch! skipped: gem not available")
+            return false
+          end
+          unless loaded?
+            Braintrust::Log.debug("#{integration_name}.patch! skipped: library not loaded")
+            return false
+          end
+          unless compatible?
+            Braintrust::Log.debug("#{integration_name}.patch! skipped: version not compatible")
+            return false
+          end
+
+          # Try all applicable patchers
+          success = false
+          patchers.each do |patch|
+            # Check if this patcher is applicable
+            next unless patch.applicable?
+
+            # Attempt to patch (patcher checks applicable? again under lock)
+            success = true if patch.patch!(**options)
+          end
+
+          Braintrust::Log.debug("#{integration_name}.patch! skipped: no applicable patcher") unless success
+          success
+        end
+
+        # Register this integration with the global registry.
+        def register!
+          Registry.instance.register(self)
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,22 @@
+# frozen_string_literal: true
+
+# Backward compatibility shim for the old OpenAI integration API.
+# This file now just delegates to the new API.
+
+module Braintrust
+  module Trace
+    module OpenAI
+      # Wrap an OpenAI::Client to automatically create spans for chat completions and responses.
+      # This is the legacy API - delegates to the new contrib framework.
+      #
+      # @param client [OpenAI::Client] the OpenAI client to wrap
+      # @param tracer_provider [OpenTelemetry::SDK::Trace::TracerProvider] the tracer provider (defaults to global)
+      # @return [OpenAI::Client] the wrapped client
+      def self.wrap(client, tracer_provider: nil)
+        Log.warn("Braintrust::Trace::OpenAI.wrap() is deprecated and will be removed in a future version: use Braintrust.instrument!() instead.")
+        Braintrust.instrument!(:openai, target: client, tracer_provider: tracer_provider)
+        client
+      end
+    end
+  end
+end

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