braintrust 0.0.12 → 0.1.0

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

@@ -0,0 +1,53 @@
+# frozen_string_literal: true
+
+module Braintrust
+  module Contrib
+    module Anthropic
+      module Instrumentation
+        # Common utilities for Anthropic SDK instrumentation.
+        module Common
+          # Parse Anthropic SDK usage tokens into normalized Braintrust metrics.
+          # Accumulates cache tokens into prompt_tokens and calculates total.
+          # Works with both Hash objects and SDK response objects (via to_h).
+          # @param usage [Hash, Object] usage object from Anthropic response
+          # @return [Hash<String, Integer>] normalized metrics
+          def self.parse_usage_tokens(usage)
+            metrics = {}
+            return metrics unless usage
+
+            usage_hash = usage.respond_to?(:to_h) ? usage.to_h : usage
+            return metrics unless usage_hash.is_a?(Hash)
+
+            # Anthropic SDK field mappings → Braintrust metrics
+            field_map = {
+              "input_tokens" => "prompt_tokens",
+              "output_tokens" => "completion_tokens",
+              "cache_read_input_tokens" => "prompt_cached_tokens",
+              "cache_creation_input_tokens" => "prompt_cache_creation_tokens"
+            }
+
+            usage_hash.each do |key, value|
+              next unless value.is_a?(Numeric)
+              key_str = key.to_s
+              target = field_map[key_str]
+              metrics[target] = value.to_i if target
+            end
+
+            # Accumulate cache tokens into prompt_tokens (matching TS/Python SDKs)
+            prompt_tokens = (metrics["prompt_tokens"] || 0) +
+              (metrics["prompt_cached_tokens"] || 0) +
+              (metrics["prompt_cache_creation_tokens"] || 0)
+            metrics["prompt_tokens"] = prompt_tokens if prompt_tokens > 0
+
+            # Calculate total
+            if metrics.key?("prompt_tokens") && metrics.key?("completion_tokens")
+              metrics["tokens"] = metrics["prompt_tokens"] + metrics["completion_tokens"]
+            end
+
+            metrics
+          end
+        end
+      end
+    end
+  end
+end

data/lib/braintrust/contrib/anthropic/instrumentation/messages.rb

@@ -0,0 +1,232 @@
+# frozen_string_literal: true
+
+require "opentelemetry/sdk"
+require "json"
+require_relative "../../support/otel"
+require_relative "common"
+require_relative "../../../internal/time"
+
+module Braintrust
+  module Contrib
+    module Anthropic
+      module Instrumentation
+        # Messages instrumentation for Anthropic.
+        # Wraps create() and stream() methods to create spans.
+        module Messages
+          def self.included(base)
+            base.prepend(InstanceMethods) unless applied?(base)
+          end
+
+          def self.applied?(base)
+            base.ancestors.include?(InstanceMethods)
+          end
+
+          module InstanceMethods
+            METADATA_FIELDS = %i[
+              model max_tokens temperature top_p top_k stop_sequences
+              stream tools tool_choice thinking metadata service_tier
+            ].freeze
+
+            # Wrap synchronous messages.create
+            def create(**params)
+              client = instance_variable_get(:@client)
+              tracer = Braintrust::Contrib.tracer_for(client)
+
+              tracer.in_span("anthropic.messages.create") do |span|
+                metadata = build_metadata(params)
+                set_input(span, params)
+
+                response = nil
+                time_to_first_token = Braintrust::Internal::Time.measure do
+                  response = super(**params)
+                end
+
+                set_output(span, response)
+                set_metrics(span, response, time_to_first_token)
+                finalize_metadata(span, metadata, response)
+
+                response
+              end
+            end
+
+            # Wrap streaming messages.stream
+            # 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,
+                messages_instance: self,
+                start_time: Braintrust::Internal::Time.measure)
+              stream_obj
+            end
+
+            private
+
+            def finalize_stream_span(span, stream_obj, metadata, time_to_first_token)
+              if stream_obj.respond_to?(:accumulated_message)
+                begin
+                  msg = stream_obj.accumulated_message
+                  set_output(span, msg)
+                  set_metrics(span, msg, time_to_first_token)
+                  metadata["stop_reason"] = msg.stop_reason if msg.respond_to?(:stop_reason) && msg.stop_reason
+                  metadata["model"] = msg.model if msg.respond_to?(:model) && msg.model
+                rescue => e
+                  Braintrust::Log.debug("Failed to get accumulated message: #{e.message}")
+                end
+              end
+              Support::OTel.set_json_attr(span, "braintrust.metadata", metadata)
+            end
+
+            def build_metadata(params, stream: false)
+              metadata = {
+                "provider" => "anthropic",
+                "endpoint" => "/v1/messages"
+              }
+              metadata["stream"] = true if stream
+              METADATA_FIELDS.each do |field|
+                metadata[field.to_s] = params[field] if params.key?(field)
+              end
+              metadata
+            end
+
+            def set_input(span, params)
+              input_messages = []
+
+              if params[:system]
+                system_content = params[:system]
+                if system_content.is_a?(Array)
+                  system_text = system_content.map { |blk|
+                    blk.is_a?(Hash) ? blk[:text] : blk
+                  }.join("\n")
+                  input_messages << {role: "system", content: system_text}
+                else
+                  input_messages << {role: "system", content: system_content}
+                end
+              end
+
+              if params[:messages]
+                messages_array = params[:messages].map(&:to_h)
+                input_messages.concat(messages_array)
+              end
+
+              Support::OTel.set_json_attr(span, "braintrust.input_json", input_messages) if input_messages.any?
+            end
+
+            def set_output(span, response)
+              return unless response.respond_to?(:content) && response.content
+
+              content_array = response.content.map(&:to_h)
+              output = [{
+                role: response.respond_to?(:role) ? response.role : "assistant",
+                content: content_array
+              }]
+              Support::OTel.set_json_attr(span, "braintrust.output_json", output)
+            end
+
+            def set_metrics(span, response, time_to_first_token)
+              metrics = {}
+              if response.respond_to?(:usage) && response.usage
+                metrics = Common.parse_usage_tokens(response.usage)
+              end
+              metrics["time_to_first_token"] = time_to_first_token if time_to_first_token
+              Support::OTel.set_json_attr(span, "braintrust.metrics", metrics) unless metrics.empty?
+            end
+
+            def finalize_metadata(span, metadata, response)
+              metadata["stop_reason"] = response.stop_reason if response.respond_to?(:stop_reason) && response.stop_reason
+              metadata["stop_sequence"] = response.stop_sequence if response.respond_to?(:stop_sequence) && response.stop_sequence
+              metadata["model"] = response.model if response.respond_to?(:model) && response.model
+              Support::OTel.set_json_attr(span, "braintrust.metadata", metadata)
+            end
+          end
+        end
+
+        # MessageStream instrumentation for Anthropic.
+        # Prepended to Anthropic::Helpers::Streaming::MessageStream to create spans on consumption.
+        module MessageStream
+          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) do
+                super do |*args|
+                  ctx[:time_to_first_token] ||= Braintrust::Internal::Time.measure(ctx[:start_time])
+                  block.call(*args)
+                end
+              end
+            end
+
+            def text
+              ctx = Braintrust::Contrib::Context.from(self)
+              return super unless ctx&.[](:tracer) && !ctx[:consumed]
+
+              original_text_enum = super
+              Enumerator.new do |output|
+                trace_consumption(ctx) do
+                  original_text_enum.each do |text_chunk|
+                    ctx[:time_to_first_token] ||= Braintrust::Internal::Time.measure(ctx[:start_time])
+                    output << text_chunk
+                  end
+                end
+              end
+            end
+
+            def close
+              ctx = Braintrust::Contrib::Context.from(self)
+              if ctx&.[](:tracer) && !ctx[:consumed]
+                # Stream closed without consumption - create minimal span
+                ctx[:consumed] = true
+                tracer = ctx[:tracer]
+                params = ctx[:params]
+                metadata = ctx[:metadata]
+                messages_instance = ctx[:messages_instance]
+
+                tracer.in_span("anthropic.messages.create") do |span|
+                  messages_instance.send(:set_input, span, params)
+                  Support::OTel.set_json_attr(span, "braintrust.metadata", metadata)
+                end
+              end
+              super
+            end
+
+            private
+
+            def trace_consumption(ctx)
+              # Mark as consumed to prevent re-entry (accumulated_message calls each internally)
+              ctx[:consumed] = true
+
+              tracer = ctx[:tracer]
+              params = ctx[:params]
+              metadata = ctx[:metadata]
+              messages_instance = ctx[:messages_instance]
+
+              tracer.in_span("anthropic.messages.create") do |span|
+                messages_instance.send(:set_input, span, params)
+                Support::OTel.set_json_attr(span, "braintrust.metadata", metadata)
+
+                yield
+
+                messages_instance.send(:finalize_stream_span, span, self, metadata, ctx[:time_to_first_token])
+              end
+            end
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,53 @@
+# frozen_string_literal: true
+
+require_relative "../integration"
+
+module Braintrust
+  module Contrib
+    module Anthropic
+      # Anthropic integration for automatic instrumentation.
+      # Instruments the anthropic gem (https://github.com/anthropics/anthropic-sdk-ruby).
+      class Integration
+        include Braintrust::Contrib::Integration
+
+        MINIMUM_VERSION = "0.3.0"
+
+        GEM_NAMES = ["anthropic"].freeze
+        REQUIRE_PATHS = ["anthropic"].freeze
+
+        # @return [Symbol] Unique identifier for this integration
+        def self.integration_name
+          :anthropic
+        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 anthropic gem is available
+        def self.loaded?
+          defined?(::Anthropic::Client) ? 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"
+          [MessagesPatcher]
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,62 @@
+# frozen_string_literal: true
+
+require_relative "../patcher"
+require_relative "instrumentation/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
+    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