braintrust 0.0.12 → 0.1.0

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

@@ -0,0 +1,464 @@
+# frozen_string_literal: true
+
+require "opentelemetry/sdk"
+require "json"
+
+require_relative "common"
+require_relative "../../support/otel"
+require_relative "../../../internal/encoding"
+require_relative "../../../internal/time"
+
+module Braintrust
+  module Contrib
+    module RubyLLM
+      module Instrumentation
+        # Chat instrumentation for RubyLLM.
+        # Wraps complete() and execute_tool() methods to create spans.
+        module Chat
+          def self.included(base)
+            # Guard against double-wrapping
+            base.prepend(InstanceMethods) unless applied?(base)
+          end
+
+          def self.applied?(base)
+            base.ancestors.include?(InstanceMethods)
+          end
+
+          module InstanceMethods
+            # Wrap complete() to trace chat completions.
+            # Each call creates a span - recursive calls from tool execution
+            # create nested spans (each is a separate API call).
+            def complete(&block)
+              return block ? super : super() unless tracing_enabled?
+
+              tracer = Braintrust::Contrib.tracer_for(self)
+
+              tracer.in_span("ruby_llm.chat") do |span|
+                if block
+                  # Streaming: pass a block that calls super() with the wrapper
+                  handle_streaming_complete(span, block) do |&wrapper|
+                    super(&wrapper)
+                  end
+                else
+                  # Non-streaming: pass a block that calls super()
+                  handle_non_streaming_complete(span) do
+                    super()
+                  end
+                end
+              end
+            end
+
+            private
+
+            # Wrap execute_tool() to trace tool executions.
+            # This is a private method in RubyLLM - wrapping it avoids
+            # conflicting with user-registered on_tool_call/on_tool_result callbacks.
+            def execute_tool(tool_call)
+              return super unless tracing_enabled?
+
+              tracer = Braintrust::Contrib.tracer_for(self)
+
+              tracer.in_span("ruby_llm.tool.#{tool_call.name}") do |span|
+                Support::OTel.set_json_attr(span, "braintrust.span_attributes", {type: "tool"})
+                span.set_attribute("tool.name", tool_call.name)
+                span.set_attribute("tool.call_id", tool_call.id)
+
+                Support::OTel.set_json_attr(span, "braintrust.input_json", {
+                  "name" => tool_call.name,
+                  "arguments" => tool_call.arguments
+                })
+
+                result = super
+
+                Support::OTel.set_json_attr(span, "braintrust.output_json", result)
+                result
+              end
+            end
+
+            # DEPRECATED: Support legacy unwrap()
+            # Checks Context for enabled: false on instance or class.
+            # This will be removed in a future version.
+            def tracing_enabled?
+              ctx = Braintrust::Contrib.context_for(self)
+              class_ctx = Braintrust::Contrib.context_for(self.class)
+              ctx&.[](:enabled) != false && class_ctx&.[](:enabled) != false
+            end
+
+            # Handle streaming complete request with tracing.
+            # Calls the provided block with a wrapper that aggregates chunks.
+            # @param span [OpenTelemetry::Trace::Span] the span to record to
+            # @param user_block [Proc] the streaming block from user
+            # @param super_caller [Proc] block that calls super(&wrapper)
+            def handle_streaming_complete(span, user_block, &super_caller)
+              aggregated_chunks = []
+              metadata = extract_metadata(stream: true)
+              input_messages = build_input_messages
+              start_time = Braintrust::Internal::Time.measure
+              time_to_first_token = nil
+
+              Support::OTel.set_json_attr(span, "braintrust.input_json", input_messages) if input_messages.any?
+              Support::OTel.set_json_attr(span, "braintrust.metadata", metadata)
+
+              # Wrapper block that RubyLLM calls once per chunk.
+              # Aggregates chunks for span recording and forwards to user's block.
+              wrapper = proc do |chunk|
+                time_to_first_token ||= Braintrust::Internal::Time.measure(start_time)
+                aggregated_chunks << chunk
+                user_block.call(chunk)
+              end
+
+              begin
+                result = super_caller.call(&wrapper)
+
+                capture_streaming_output(span, aggregated_chunks, result, time_to_first_token)
+                result
+              rescue => e
+                span.record_exception(e)
+                span.status = ::OpenTelemetry::Trace::Status.error("RubyLLM error: #{e.message}")
+                raise
+              end
+            end
+
+            # Handle non-streaming complete request with tracing.
+            # Calls the provided block to invoke super() and returns the response.
+            # @param span [OpenTelemetry::Trace::Span] the span to record to
+            # @param super_caller [Proc] block that calls super()
+            def handle_non_streaming_complete(span, &super_caller)
+              metadata = extract_metadata
+              input_messages = build_input_messages
+              Support::OTel.set_json_attr(span, "braintrust.input_json", input_messages) if input_messages.any?
+
+              messages_before_count = messages&.length || 0
+
+              begin
+                response = nil
+                time_to_first_token = Braintrust::Internal::Time.measure do
+                  response = super_caller.call
+                end
+
+                capture_non_streaming_output(span, response, messages_before_count, time_to_first_token)
+                Support::OTel.set_json_attr(span, "braintrust.metadata", metadata)
+
+                response
+              rescue => e
+                span.record_exception(e)
+                span.status = ::OpenTelemetry::Trace::Status.error("RubyLLM error: #{e.message}")
+                raise
+              end
+            end
+
+            # Extract metadata from chat instance (provider, model, tools, stream flag)
+            def extract_metadata(stream: false)
+              metadata = {"provider" => "ruby_llm"}
+              metadata["stream"] = true if stream
+
+              # Extract model
+              if respond_to?(:model) && model
+                model_id = model.respond_to?(:id) ? model.id : model.to_s
+                metadata["model"] = model_id
+              end
+
+              # Extract tools (only for non-streaming)
+              if !stream && respond_to?(:tools) && tools&.any?
+                metadata["tools"] = extract_tools_metadata
+              end
+
+              metadata
+            end
+
+            # Extract tools metadata from chat instance
+            def extract_tools_metadata
+              provider = instance_variable_get(:@provider) if instance_variable_defined?(:@provider)
+
+              tools.map do |_name, tool|
+                format_tool_schema(tool, provider)
+              end
+            end
+
+            # Format a tool into OpenAI-compatible schema
+            def format_tool_schema(tool, provider)
+              tool_schema = nil
+
+              # Use provider-specific tool_for method if available
+              if provider
+                begin
+                  tool_schema = if provider.is_a?(::RubyLLM::Providers::OpenAI)
+                    ::RubyLLM::Providers::OpenAI::Tools.tool_for(tool)
+                  elsif defined?(::RubyLLM::Providers::Anthropic) && provider.is_a?(::RubyLLM::Providers::Anthropic)
+                    ::RubyLLM::Providers::Anthropic::Tools.tool_for(tool)
+                  elsif tool.respond_to?(:params_schema) && tool.params_schema
+                    build_basic_tool_schema(tool)
+                  else
+                    build_minimal_tool_schema(tool)
+                  end
+                rescue NameError, ArgumentError => e
+                  Braintrust::Log.debug("Failed to extract tool schema using provider-specific method: #{e.class.name}: #{e.message}")
+                  tool_schema = (tool.respond_to?(:params_schema) && tool.params_schema) ? build_basic_tool_schema(tool) : build_minimal_tool_schema(tool)
+                end
+              else
+                tool_schema = (tool.respond_to?(:params_schema) && tool.params_schema) ? build_basic_tool_schema(tool) : build_minimal_tool_schema(tool)
+              end
+
+              # Strip RubyLLM-specific fields to match native OpenAI format
+              function_key = tool_schema&.key?(:function) ? :function : "function"
+              if tool_schema && tool_schema[function_key]
+                tool_params = tool_schema[function_key][:parameters] || tool_schema[function_key]["parameters"]
+                if tool_params.is_a?(Hash)
+                  tool_params = tool_params.dup if tool_params.frozen?
+                  tool_params.delete("strict")
+                  tool_params.delete(:strict)
+                  tool_params.delete("additionalProperties")
+                  tool_params.delete(:additionalProperties)
+                  params_key = tool_schema[function_key].key?(:parameters) ? :parameters : "parameters"
+                  tool_schema[function_key][params_key] = tool_params
+                end
+              end
+
+              tool_schema
+            end
+
+            # Build a basic tool schema with parameters
+            def build_basic_tool_schema(tool)
+              {
+                "type" => "function",
+                "function" => {
+                  "name" => tool.name.to_s,
+                  "description" => tool.description,
+                  "parameters" => tool.params_schema
+                }
+              }
+            end
+
+            # Build a minimal tool schema without parameters
+            def build_minimal_tool_schema(tool)
+              {
+                "type" => "function",
+                "function" => {
+                  "name" => tool.name.to_s,
+                  "description" => tool.description,
+                  "parameters" => {}
+                }
+              }
+            end
+
+            # Build input messages array from chat history
+            def build_input_messages
+              return [] unless respond_to?(:messages) && messages&.any?
+
+              messages.map { |m| format_message_for_input(m) }
+            end
+
+            # Format a RubyLLM message to OpenAI-compatible format
+            def format_message_for_input(msg)
+              formatted = {"role" => msg.role.to_s}
+
+              # Handle content
+              if msg.respond_to?(:content) && msg.content
+                raw_content = msg.content
+
+                # Check if content is a Content object with attachments (issue #71)
+                formatted["content"] = if raw_content.respond_to?(:text) && raw_content.respond_to?(:attachments) && raw_content.attachments&.any?
+                  format_multipart_content(raw_content)
+                else
+                  format_simple_content(raw_content, msg.role.to_s)
+                end
+              end
+
+              # Handle tool_calls for assistant messages
+              if msg.respond_to?(:tool_calls) && msg.tool_calls&.any?
+                formatted["tool_calls"] = format_tool_calls(msg.tool_calls)
+                formatted["content"] = nil
+              end
+
+              # Handle tool_call_id for tool result messages
+              if msg.respond_to?(:tool_call_id) && msg.tool_call_id
+                formatted["tool_call_id"] = msg.tool_call_id
+              end
+
+              formatted
+            end
+
+            # Format multipart content with text and attachments
+            # @param content_obj [Object] Content object with text and attachments
+            # @return [Array<Hash>] array of content parts
+            def format_multipart_content(content_obj)
+              content_parts = []
+
+              # Add text part
+              content_parts << {"type" => "text", "text" => content_obj.text} if content_obj.text
+
+              # Add attachment parts (convert to Braintrust format)
+              content_obj.attachments.each do |attachment|
+                content_parts << format_attachment_for_input(attachment)
+              end
+
+              content_parts
+            end
+
+            # Format simple text content
+            # @param raw_content [Object] String or Content object with text
+            # @param role [String] the message role
+            # @return [String] formatted text content
+            def format_simple_content(raw_content, role)
+              content = raw_content
+              content = content.text if content.respond_to?(:text)
+
+              # Convert Ruby hash string to JSON for tool results
+              if role == "tool" && content.is_a?(String) && content.start_with?("{:")
+                begin
+                  content = content.gsub(/(?<=\{|, ):(\w+)=>/, '"\1":').gsub("=>", ":")
+                rescue
+                  # Keep original if conversion fails
+                end
+              end
+
+              content
+            end
+
+            # Format a RubyLLM attachment to OpenAI-compatible format
+            # @param attachment [Object] the RubyLLM attachment
+            # @return [Hash] OpenAI image_url format for consistency with other integrations
+            def format_attachment_for_input(attachment)
+              # RubyLLM Attachment has: source (Pathname), filename, mime_type
+              if attachment.respond_to?(:source) && attachment.source
+                begin
+                  data = File.binread(attachment.source.to_s)
+                  encoded = Braintrust::Internal::Encoding::Base64.strict_encode64(data)
+                  mime_type = attachment.respond_to?(:mime_type) ? attachment.mime_type : "application/octet-stream"
+
+                  # Use OpenAI's image_url format for consistency
+                  {
+                    "type" => "image_url",
+                    "image_url" => {
+                      "url" => "data:#{mime_type};base64,#{encoded}"
+                    }
+                  }
+                rescue => e
+                  Braintrust::Log.debug("Failed to read attachment file: #{e.message}")
+                  # Return a placeholder if we can't read the file
+                  {"type" => "text", "text" => "[attachment: #{attachment.respond_to?(:filename) ? attachment.filename : "unknown"}]"}
+                end
+              elsif attachment.respond_to?(:to_h)
+                # Try to use attachment's own serialization
+                attachment.to_h
+              else
+                {"type" => "text", "text" => "[attachment]"}
+              end
+            end
+
+            # Format tool calls into OpenAI format
+            def format_tool_calls(tool_calls)
+              tool_calls.map do |_id, tc|
+                args = tc.arguments
+                args_string = args.is_a?(String) ? args : JSON.generate(args)
+
+                {
+                  "id" => tc.id,
+                  "type" => "function",
+                  "function" => {
+                    "name" => tc.name,
+                    "arguments" => args_string
+                  }
+                }
+              end
+            end
+
+            # Capture streaming output and metrics
+            def capture_streaming_output(span, aggregated_chunks, result, time_to_first_token)
+              return if aggregated_chunks.empty?
+
+              # Aggregate content from chunks
+              # Extract text from Content objects if present (issue #71)
+              aggregated_content = aggregated_chunks.map { |c|
+                content = c.respond_to?(:content) ? c.content : c.to_s
+                content = content.text if content.respond_to?(:text)
+                content
+              }.join
+
+              output = [{
+                role: "assistant",
+                content: aggregated_content
+              }]
+              Support::OTel.set_json_attr(span, "braintrust.output_json", output)
+
+              # Set metrics (token usage + time_to_first_token)
+              # RubyLLM stores tokens directly in the response hash, not in a nested usage object
+              metrics = {}
+              if result.respond_to?(:to_h)
+                result_hash = result.to_h
+                usage = {
+                  "input_tokens" => result_hash[:input_tokens],
+                  "output_tokens" => result_hash[:output_tokens],
+                  "cached_tokens" => result_hash[:cached_tokens],
+                  "cache_creation_tokens" => result_hash[:cache_creation_tokens]
+                }.compact
+
+                unless usage.empty?
+                  metrics = Common.parse_usage_tokens(usage)
+                end
+              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
+
+            # Capture non-streaming output and metrics
+            def capture_non_streaming_output(span, response, messages_before_count, time_to_first_token)
+              return unless response
+
+              message = {
+                "role" => "assistant",
+                "content" => nil
+              }
+
+              # Add content if it's a simple text response
+              # Extract text from Content objects if present (issue #71)
+              if response.respond_to?(:content) && response.content && !response.content.empty?
+                content = response.content
+                content = content.text if content.respond_to?(:text)
+                message["content"] = content
+              end
+
+              # Check if there are tool calls in the messages history
+              if respond_to?(:messages) && messages
+                assistant_msg = messages[messages_before_count..]&.find { |m|
+                  m.role.to_s == "assistant" && m.respond_to?(:tool_calls) && m.tool_calls&.any?
+                }
+
+                if assistant_msg&.tool_calls&.any?
+                  message["tool_calls"] = format_tool_calls(assistant_msg.tool_calls)
+                  message["content"] = nil
+                end
+              end
+
+              output = [{
+                "index" => 0,
+                "message" => message,
+                "finish_reason" => message["tool_calls"] ? "tool_calls" : "stop"
+              }]
+
+              Support::OTel.set_json_attr(span, "braintrust.output_json", output)
+
+              # Set metrics (token usage + time_to_first_token)
+              metrics = {}
+              if response.respond_to?(:to_h)
+                response_hash = response.to_h
+                usage = {
+                  "input_tokens" => response_hash[:input_tokens],
+                  "output_tokens" => response_hash[:output_tokens],
+                  "cached_tokens" => response_hash[:cached_tokens],
+                  "cache_creation_tokens" => response_hash[:cache_creation_tokens]
+                }.compact
+
+                unless usage.empty?
+                  metrics = Common.parse_usage_tokens(usage)
+                end
+              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
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,58 @@
+# frozen_string_literal: true
+
+module Braintrust
+  module Contrib
+    module RubyLLM
+      module Instrumentation
+        # Common utilities for RubyLLM instrumentation.
+        module Common
+          # Parse RubyLLM usage tokens into normalized Braintrust metrics.
+          # RubyLLM normalizes token fields from all providers (OpenAI, Anthropic, etc.)
+          # into a consistent format:
+          #   - input_tokens: prompt tokens sent
+          #   - output_tokens: completion tokens received
+          #   - cached_tokens: tokens read from cache
+          #   - cache_creation_tokens: tokens written to cache
+          #
+          # @param usage [Hash, Object] usage object from RubyLLM response
+          # @return [Hash<String, Integer>] normalized metrics for Braintrust
+          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)
+
+            # RubyLLM normalized field mappings → Braintrust metrics
+            field_map = {
+              "input_tokens" => "prompt_tokens",
+              "output_tokens" => "completion_tokens",
+              "cached_tokens" => "prompt_cached_tokens",
+              "cache_creation_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/ruby_llm/integration.rb

@@ -0,0 +1,54 @@
+# frozen_string_literal: true
+
+require_relative "../integration"
+require_relative "deprecated"
+
+module Braintrust
+  module Contrib
+    module RubyLLM
+      # RubyLLM integration for automatic instrumentation.
+      # Instruments the crmne/ruby_llm gem.
+      class Integration
+        include Braintrust::Contrib::Integration
+
+        MINIMUM_VERSION = "1.8.0"
+
+        GEM_NAMES = ["ruby_llm"].freeze
+        REQUIRE_PATHS = ["ruby_llm"].freeze
+
+        # @return [Symbol] Unique identifier for this integration
+        def self.integration_name
+          :ruby_llm
+        end
+
+        # @return [Array<String>] Gem names this integration supports
+        def self.gem_names
+          GEM_NAMES
+        end
+
+        # @return [Array<String>] Require paths for auto-instrument detection
+        def self.require_paths
+          REQUIRE_PATHS
+        end
+
+        # @return [String] Minimum compatible version
+        def self.minimum_version
+          MINIMUM_VERSION
+        end
+
+        # @return [Boolean] true if ruby_llm gem is available
+        def self.loaded?
+          defined?(::RubyLLM::Chat) ? true : false
+        end
+
+        # Lazy-load the patcher only when actually patching.
+        # This keeps the integration stub lightweight.
+        # @return [Array<Class>] The patcher classes
+        def self.patchers
+          require_relative "patcher"
+          [ChatPatcher]
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,44 @@
+# frozen_string_literal: true
+
+require_relative "../patcher"
+require_relative "instrumentation/chat"
+
+module Braintrust
+  module Contrib
+    module RubyLLM
+      # Patcher for RubyLLM chat completions.
+      # Instruments RubyLLM::Chat#complete and #execute_tool methods.
+      class ChatPatcher < Braintrust::Contrib::Patcher
+        class << self
+          def applicable?
+            defined?(::RubyLLM::Chat)
+          end
+
+          def patched?(**options)
+            target_class = options[:target]&.singleton_class || ::RubyLLM::Chat
+            Instrumentation::Chat.applied?(target_class)
+          end
+
+          # Perform the actual patching.
+          # @param options [Hash] Configuration options passed from integration
+          # @option options [Object] :target Optional target instance to patch
+          # @option options [OpenTelemetry::SDK::Trace::TracerProvider] :tracer_provider Optional tracer provider
+          # @return [void]
+          def perform_patch(**options)
+            return unless applicable?
+
+            if options[:target]
+              # Instance-level (for only this chat instance)
+              raise ArgumentError, "target must be a kind of ::RubyLLM::Chat" unless options[:target].is_a?(::RubyLLM::Chat)
+
+              options[:target].singleton_class.include(Instrumentation::Chat)
+            else
+              # Class-level (for all chat instances)
+              ::RubyLLM::Chat.include(Instrumentation::Chat)
+            end
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,24 @@
+# frozen_string_literal: true
+
+# Backward compatibility shim for the old ruby-openai integration API.
+# This file now just delegates to the new API.
+
+module Braintrust
+  module Trace
+    module AlexRudall
+      module RubyOpenAI
+        # Wrap an OpenAI::Client (ruby-openai gem) to automatically create spans.
+        # 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::AlexRudall::RubyOpenAI.wrap() is deprecated and will be removed in a future version: use Braintrust.instrument!() instead.")
+          Braintrust.instrument!(:ruby_openai, target: client, tracer_provider: tracer_provider)
+          client
+        end
+      end
+    end
+  end
+end