braintrust 0.0.7 → 0.0.8

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: ad055b60c4efb984bce955b1c00c684c760c849dea7a54d49a452f925dbab629
-  data.tar.gz: fe271abfac7810e53ff88efb139bc41b519799ed7dbd196ab5bb64fcbc35b62c
+  metadata.gz: c6b2bcda06084f2e90d2602659ca71cf0ab574ac8c74c367890cbb2b04740529
+  data.tar.gz: 306b5a46660eae3d3e3811d021627883419a4dc4c114e51e40be64c590868c95
 SHA512:
-  metadata.gz: a33fc58073542bf7d7dbf45092d9f3a6669d7f13fa98d3d8753d8fccd456c7e010b8a3ac0be5625e6761e164d2223f20b7fbaa3f02ead978f47398153c6c8ac2
-  data.tar.gz: b24f1377f4ec25f09c7c5e1366e1428c9164d7b3b7bdcf82331aa84d10603b250c61983ea811a13e7be1a1276558ef2995292fbefb20e99a374e59e9eaefb8b1
+  metadata.gz: 1db7bf706b260762aa114eb5e8f844cb0567efd5a6f9d8cca03667111c0e89ff68f4e53b3a3adc6ad2192947602fc4b88a8e0057169ad7eff12ccb1c2ecb4951
+  data.tar.gz: bbc71c33bb28da124bd1cc61c8bf4f765ec2899a57bf604da624a04c42a7bfce508ed1eb78c4ec421da5038fbaf89086cee8d0b04998d223568df05e8640679f

data/lib/braintrust/trace/contrib/github.com/alexrudall/ruby-openai/ruby-openai.rb

@@ -28,6 +28,105 @@ module Braintrust
               Braintrust::Trace.parse_openai_usage_tokens(usage)
             end
 
+            # Aggregate streaming chunks into a single response structure
+            # @param chunks [Array<Hash>] array of chunk hashes from stream
+            # @return [Hash] aggregated response with choices, usage, id, created, model
+            def self.aggregate_streaming_chunks(chunks)
+              return {} if chunks.empty?
+
+              # Initialize aggregated structure
+              aggregated = {
+                "id" => nil,
+                "created" => nil,
+                "model" => nil,
+                "usage" => nil,
+                "choices" => []
+              }
+
+              # Track aggregated content for the first choice
+              role = nil
+              content = +""
+
+              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"]
+
+                # Aggregate usage (usually only in last chunk if stream_options.include_usage is set)
+                aggregated["usage"] = chunk["usage"] if chunk["usage"]
+
+                # Aggregate content from first choice
+                if chunk.dig("choices", 0, "delta", "role")
+                  role ||= chunk.dig("choices", 0, "delta", "role")
+                end
+                if chunk.dig("choices", 0, "delta", "content")
+                  content << chunk.dig("choices", 0, "delta", "content")
+                end
+              end
+
+              # Build aggregated choices array
+              aggregated["choices"] = [
+                {
+                  "index" => 0,
+                  "message" => {
+                    "role" => role || "assistant",
+                    "content" => content
+                  },
+                  "finish_reason" => chunks.dig(-1, "choices", 0, "finish_reason")
+                }
+              ]
+
+              aggregated
+            end
+
+            # Aggregate responses streaming chunks into a single response structure
+            # @param chunks [Array<Hash>] array of chunk hashes from stream
+            # @return [Hash] aggregated response with output, usage, id
+            def self.aggregate_responses_chunks(chunks)
+              return {} if chunks.empty?
+
+              # Find the response.completed event which has the final response
+              completed_chunk = chunks.find { |c| c["type"] == "response.completed" }
+
+              if completed_chunk && completed_chunk["response"]
+                response = completed_chunk["response"]
+                return {
+                  "id" => response["id"],
+                  "output" => response["output"],
+                  "usage" => response["usage"]
+                }
+              end
+
+              # Fallback if no completed event found
+              {}
+            end
+
+            # Set span attributes from response data (works for both streaming and non-streaming)
+            # @param span [OpenTelemetry::Trace::Span] the span to set attributes on
+            # @param response_data [Hash] response hash with keys: choices, usage, id, created, model, system_fingerprint, service_tier
+            # @param time_to_first_token [Float] time to first token in seconds
+            # @param metadata [Hash] metadata hash to update with response fields
+            def self.set_span_attributes(span, response_data, time_to_first_token, metadata)
+              # Set output (choices) as JSON
+              if response_data["choices"]&.any?
+                set_json_attr(span, "braintrust.output_json", response_data["choices"])
+              end
+
+              # Set metrics (token usage + time_to_first_token)
+              metrics = {}
+              if response_data["usage"]
+                metrics = parse_usage_tokens(response_data["usage"])
+              end
+              metrics["time_to_first_token"] = time_to_first_token || 0.0
+              set_json_attr(span, "braintrust.metrics", metrics) unless metrics.empty?
+
+              # Update metadata with response fields
+              %w[id created model system_fingerprint service_tier].each do |field|
+                metadata[field] = response_data[field] if response_data[field]
+              end
+            end
+
             # Wrap an OpenAI::Client (ruby-openai gem) to automatically create spans
             # Supports both synchronous and streaming requests
             # @param client [OpenAI::Client] the OpenAI client to wrap
@@ -35,92 +134,235 @@ module Braintrust
             def self.wrap(client, tracer_provider: nil)
               tracer_provider ||= ::OpenTelemetry.tracer_provider
 
+              # Store tracer provider on the client for use by wrapper modules
+              client.instance_variable_set(:@braintrust_tracer_provider, tracer_provider)
+
               # Wrap chat completions
-              wrap_chat(client, tracer_provider)
+              wrap_chat(client)
+
+              # Wrap responses API if available
+              wrap_responses(client) if client.respond_to?(:responses)
 
               client
             end
 
             # Wrap chat API
             # @param client [OpenAI::Client] the OpenAI client
-            # @param tracer_provider [OpenTelemetry::SDK::Trace::TracerProvider] the tracer provider
-            def self.wrap_chat(client, tracer_provider)
-              # Create a wrapper module that intercepts the chat method
-              wrapper = Module.new do
-                define_method(:chat) do |parameters:|
-                  tracer = tracer_provider.tracer("braintrust")
-
-                  tracer.in_span("openai.chat.completions.create") do |span|
-                    # Initialize metadata hash
-                    metadata = {
-                      "provider" => "openai",
-                      "endpoint" => "/v1/chat/completions"
-                    }
-
-                    # Capture request metadata fields
-                    metadata_fields = %w[
-                      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
-                    ]
-
-                    metadata_fields.each do |field|
-                      field_sym = field.to_sym
-                      if parameters.key?(field_sym)
-                        # Special handling for stream parameter (it's a Proc)
-                        metadata[field] = if field == "stream"
-                          true  # Just mark as streaming
-                        else
-                          parameters[field_sym]
-                        end
+            def self.wrap_chat(client)
+              client.singleton_class.prepend(ChatWrapper)
+            end
+
+            # Wrap responses API
+            # @param client [OpenAI::Client] the OpenAI client
+            def self.wrap_responses(client)
+              # Store tracer provider on the responses object for use by wrapper module
+              responses_obj = client.responses
+              responses_obj.instance_variable_set(:@braintrust_tracer_provider, client.instance_variable_get(:@braintrust_tracer_provider))
+              responses_obj.singleton_class.prepend(ResponsesCreateWrapper)
+            end
+
+            # Wrapper module for chat completions
+            module ChatWrapper
+              def chat(parameters:)
+                tracer_provider = @braintrust_tracer_provider
+                tracer = tracer_provider.tracer("braintrust")
+
+                tracer.in_span("Chat Completion") do |span|
+                  # Track start time for time_to_first_token
+                  start_time = Time.now
+                  time_to_first_token = nil
+                  is_streaming = parameters.key?(:stream) && parameters[:stream].is_a?(Proc)
+
+                  # Initialize metadata hash
+                  metadata = {
+                    "provider" => "openai",
+                    "endpoint" => "/v1/chat/completions"
+                  }
+
+                  # Capture request metadata fields
+                  metadata_fields = %w[
+                    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
+                  ]
+
+                  metadata_fields.each do |field|
+                    field_sym = field.to_sym
+                    if parameters.key?(field_sym)
+                      # Special handling for stream parameter (it's a Proc)
+                      metadata[field] = if field == "stream"
+                        true  # Just mark as streaming
+                      else
+                        parameters[field_sym]
                       end
                     end
+                  end
+
+                  # Set input messages as JSON
+                  if parameters[:messages]
+                    RubyOpenAI.set_json_attr(span, "braintrust.input_json", parameters[:messages])
+                  end
 
-                    # Set input messages as JSON
-                    if parameters[:messages]
-                      span.set_attribute("braintrust.input_json", JSON.generate(parameters[:messages]))
+                  # Wrap streaming callback if present to capture time to first token and aggregate chunks
+                  aggregated_chunks = []
+                  if is_streaming
+                    original_stream_proc = parameters[:stream]
+                    parameters = parameters.dup
+                    parameters[:stream] = proc do |chunk, bytesize|
+                      # Capture time to first token on first chunk
+                      time_to_first_token ||= Time.now - start_time
+                      # Aggregate chunks for later processing
+                      aggregated_chunks << chunk
+                      # Call original callback
+                      original_stream_proc.call(chunk, bytesize)
                     end
+                  end
+
+                  begin
+                    # Call the original method
+                    response = super(parameters: parameters)
+
+                    # Calculate time to first token for non-streaming
+                    time_to_first_token ||= Time.now - start_time unless is_streaming
+
+                    # Process response data
+                    if is_streaming && !aggregated_chunks.empty?
+                      # Aggregate streaming chunks into response-like structure
+                      aggregated_response = RubyOpenAI.aggregate_streaming_chunks(aggregated_chunks)
+                      RubyOpenAI.set_span_attributes(span, aggregated_response, time_to_first_token, metadata)
+                    else
+                      # Non-streaming: use response object directly
+                      RubyOpenAI.set_span_attributes(span, response || {}, time_to_first_token, metadata)
+                    end
+
+                    # Set metadata ONCE at the end with complete hash
+                    RubyOpenAI.set_json_attr(span, "braintrust.metadata", metadata)
 
-                    begin
-                      # Call the original method
-                      response = super(parameters: parameters)
+                    response
+                  rescue => e
+                    # Record exception in span
+                    span.record_exception(e)
+                    span.status = OpenTelemetry::Trace::Status.error("Exception: #{e.class} - #{e.message}")
+                    raise
+                  end
+                end
+              end
+            end
+
+            # Wrapper module for responses API create method
+            module ResponsesCreateWrapper
+              def create(parameters:)
+                tracer_provider = @braintrust_tracer_provider
+                tracer = tracer_provider.tracer("braintrust")
 
-                      # Set output (choices) as JSON
-                      if response && response["choices"]&.any?
-                        span.set_attribute("braintrust.output_json", JSON.generate(response["choices"]))
+                tracer.in_span("openai.responses.create") do |span|
+                  # Track start time for time_to_first_token
+                  start_time = Time.now
+                  time_to_first_token = nil
+                  is_streaming = parameters.key?(:stream) && parameters[:stream].is_a?(Proc)
+
+                  # Initialize metadata hash
+                  metadata = {
+                    "provider" => "openai",
+                    "endpoint" => "/v1/responses"
+                  }
+
+                  # Capture request metadata fields
+                  metadata_fields = %w[
+                    model instructions modalities tools parallel_tool_calls
+                    tool_choice temperature max_tokens top_p frequency_penalty
+                    presence_penalty seed user store response_format
+                    reasoning previous_response_id truncation
+                  ]
+
+                  metadata_fields.each do |field|
+                    field_sym = field.to_sym
+                    if parameters.key?(field_sym)
+                      metadata[field] = parameters[field_sym]
+                    end
+                  end
+
+                  # Mark as streaming if applicable
+                  metadata["stream"] = true if is_streaming
+
+                  # Set input as JSON
+                  if parameters[:input]
+                    RubyOpenAI.set_json_attr(span, "braintrust.input_json", parameters[:input])
+                  end
+
+                  # Wrap streaming callback if present to capture time to first token and aggregate chunks
+                  aggregated_chunks = []
+                  if is_streaming
+                    original_stream_proc = parameters[:stream]
+                    parameters = parameters.dup
+                    parameters[:stream] = proc do |chunk, event|
+                      # Capture time to first token on first chunk
+                      time_to_first_token ||= Time.now - start_time
+                      # Aggregate chunks for later processing
+                      aggregated_chunks << chunk
+                      # Call original callback
+                      original_stream_proc.call(chunk, event)
+                    end
+                  end
+
+                  begin
+                    # Call the original method
+                    response = super(parameters: parameters)
+
+                    # Calculate time to first token for non-streaming
+                    time_to_first_token ||= Time.now - start_time unless is_streaming
+
+                    # Process response data
+                    if is_streaming && !aggregated_chunks.empty?
+                      # Aggregate streaming chunks into response-like structure
+                      aggregated_response = RubyOpenAI.aggregate_responses_chunks(aggregated_chunks)
+
+                      # Set output as JSON
+                      if aggregated_response["output"]
+                        RubyOpenAI.set_json_attr(span, "braintrust.output_json", aggregated_response["output"])
                       end
 
-                      # Set metrics (token usage)
-                      if response && response["usage"]
-                        metrics = Braintrust::Trace::Contrib::Github::Alexrudall::RubyOpenAI.parse_usage_tokens(response["usage"])
-                        span.set_attribute("braintrust.metrics", JSON.generate(metrics)) unless metrics.empty?
+                      # Set metrics (token usage + time_to_first_token)
+                      metrics = {}
+                      if aggregated_response["usage"]
+                        metrics = RubyOpenAI.parse_usage_tokens(aggregated_response["usage"])
                       end
+                      metrics["time_to_first_token"] = time_to_first_token || 0.0
+                      RubyOpenAI.set_json_attr(span, "braintrust.metrics", metrics) unless metrics.empty?
 
-                      # Add response metadata fields
-                      if response
-                        metadata["id"] = response["id"] if response["id"]
-                        metadata["created"] = response["created"] if response["created"]
-                        metadata["system_fingerprint"] = response["system_fingerprint"] if response["system_fingerprint"]
-                        metadata["service_tier"] = response["service_tier"] if response["service_tier"]
+                      # Update metadata with response fields
+                      metadata["id"] = aggregated_response["id"] if aggregated_response["id"]
+                    else
+                      # Non-streaming: use response object directly
+                      if response && response["output"]
+                        RubyOpenAI.set_json_attr(span, "braintrust.output_json", response["output"])
                       end
 
-                      # Set metadata ONCE at the end with complete hash
-                      span.set_attribute("braintrust.metadata", JSON.generate(metadata))
+                      # Set metrics (token usage + time_to_first_token)
+                      metrics = {}
+                      if response && response["usage"]
+                        metrics = RubyOpenAI.parse_usage_tokens(response["usage"])
+                      end
+                      metrics["time_to_first_token"] = time_to_first_token || 0.0
+                      RubyOpenAI.set_json_attr(span, "braintrust.metrics", metrics) unless metrics.empty?
 
-                      response
-                    rescue => e
-                      # Record exception in span
-                      span.record_exception(e)
-                      span.status = OpenTelemetry::Trace::Status.error("Exception: #{e.class} - #{e.message}")
-                      raise
+                      # Update metadata with response fields
+                      metadata["id"] = response["id"] if response && response["id"]
                     end
+
+                    # Set metadata ONCE at the end with complete hash
+                    RubyOpenAI.set_json_attr(span, "braintrust.metadata", metadata)
+
+                    response
+                  rescue => e
+                    # Record exception in span
+                    span.record_exception(e)
+                    span.status = OpenTelemetry::Trace::Status.error("Exception: #{e.class} - #{e.message}")
+                    raise
                   end
                 end
               end
-
-              # Prepend the wrapper to the client's singleton class
-              client.singleton_class.prepend(wrapper)
             end
           end
         end

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

@@ -155,7 +155,10 @@ module Braintrust
           define_method(:create) do |**params|
             tracer = tracer_provider.tracer("braintrust")
 
-            tracer.in_span("openai.chat.completions.create") do |span|
+            tracer.in_span("Chat Completion") do |span|
+              # Track start time for time_to_first_token
+              start_time = Time.now
+
               # Initialize metadata hash
               metadata = {
                 "provider" => "openai",
@@ -184,6 +187,9 @@ module Braintrust
               # Call the original method
               response = super(**params)
 
+              # Calculate time to first token
+              time_to_first_token = Time.now - start_time
+
               # Set output (choices) as JSON
               # Use to_h to get the raw structure with all fields (including tool_calls)
               if response.respond_to?(:choices) && response.choices&.any?
@@ -192,10 +198,13 @@ module Braintrust
               end
 
               # Set metrics (token usage with advanced details)
+              metrics = {}
               if response.respond_to?(:usage) && response.usage
                 metrics = Braintrust::Trace::OpenAI.parse_usage_tokens(response.usage)
-                span.set_attribute("braintrust.metrics", JSON.generate(metrics)) unless metrics.empty?
               end
+              # Add time_to_first_token metric
+              metrics["time_to_first_token"] = time_to_first_token
+              span.set_attribute("braintrust.metrics", JSON.generate(metrics)) unless metrics.empty?
 
               # Add response metadata fields
               metadata["id"] = response.id if response.respond_to?(:id) && response.id
@@ -214,13 +223,15 @@ module Braintrust
           define_method(:stream_raw) do |**params|
             tracer = tracer_provider.tracer("braintrust")
             aggregated_chunks = []
+            start_time = Time.now
+            time_to_first_token = nil
             metadata = {
               "provider" => "openai",
               "endpoint" => "/v1/chat/completions"
             }
 
             # Start span with proper context (will be child of current span if any)
-            span = tracer.start_span("openai.chat.completions.create")
+            span = tracer.start_span("Chat Completion")
 
             # Capture request metadata fields
             metadata_fields = %i[
@@ -259,6 +270,8 @@ module Braintrust
             original_each = stream.method(:each)
             stream.define_singleton_method(:each) do |&block|
               original_each.call do |chunk|
+                # Capture time to first token on first chunk
+                time_to_first_token ||= Time.now - start_time
                 aggregated_chunks << chunk.to_h
                 block&.call(chunk)
               end
@@ -275,10 +288,13 @@ module Braintrust
                 Braintrust::Trace::OpenAI.set_json_attr(span, "braintrust.output_json", aggregated_output[:choices])
 
                 # Set metrics if usage is included (requires stream_options.include_usage)
+                metrics = {}
                 if aggregated_output[:usage]
                   metrics = Braintrust::Trace::OpenAI.parse_usage_tokens(aggregated_output[:usage])
-                  Braintrust::Trace::OpenAI.set_json_attr(span, "braintrust.metrics", metrics) unless metrics.empty?
                 end
+                # Add time_to_first_token metric
+                metrics["time_to_first_token"] = time_to_first_token || 0.0
+                Braintrust::Trace::OpenAI.set_json_attr(span, "braintrust.metrics", metrics) unless metrics.empty?
 
                 # Update metadata with response fields
                 metadata["id"] = aggregated_output[:id] if aggregated_output[:id]
@@ -297,13 +313,15 @@ module Braintrust
           # Wrap stream for streaming chat completions (returns ChatCompletionStream with convenience methods)
           define_method(:stream) do |**params|
             tracer = tracer_provider.tracer("braintrust")
+            start_time = Time.now
+            time_to_first_token = nil
             metadata = {
               "provider" => "openai",
               "endpoint" => "/v1/chat/completions"
             }
 
             # Start span with proper context (will be child of current span if any)
-            span = tracer.start_span("openai.chat.completions.create")
+            span = tracer.start_span("Chat Completion")
 
             # Capture request metadata fields
             metadata_fields = %i[
@@ -354,10 +372,13 @@ module Braintrust
               end
 
               # Set metrics if usage is available
+              metrics = {}
               if snapshot.usage
                 metrics = Braintrust::Trace::OpenAI.parse_usage_tokens(snapshot.usage)
-                set_json_attr.call("braintrust.metrics", metrics) unless metrics.empty?
               end
+              # Add time_to_first_token metric
+              metrics["time_to_first_token"] = time_to_first_token || 0.0
+              set_json_attr.call("braintrust.metrics", metrics) unless metrics.empty?
 
               # Update metadata with response fields
               metadata["id"] = snapshot.id if snapshot.respond_to?(:id) && snapshot.id
@@ -378,7 +399,11 @@ module Braintrust
             # Wrap .each() method - this is the core consumption method
             original_each = stream.method(:each)
             stream.define_singleton_method(:each) do |&block|
-              original_each.call(&block)
+              original_each.call do |chunk|
+                # Capture time to first token on first chunk
+                time_to_first_token ||= Time.now - start_time
+                block&.call(chunk)
+              end
             rescue => e
               span.record_exception(e)
               span.status = ::OpenTelemetry::Trace::Status.error("Streaming error: #{e.message}")
@@ -392,8 +417,13 @@ module Braintrust
             stream.define_singleton_method(:text) do
               text_enum = original_text.call
               # Wrap the returned enumerable's .each method
+              original_text_each = text_enum.method(:each)
               text_enum.define_singleton_method(:each) do |&block|
-                super(&block)
+                original_text_each.call do |delta|
+                  # Capture time to first token on first delta
+                  time_to_first_token ||= Time.now - start_time
+                  block&.call(delta)
+                end
               rescue => e
                 span.record_exception(e)
                 span.status = ::OpenTelemetry::Trace::Status.error("Streaming error: #{e.message}")

data/lib/braintrust/trace/tokens.rb

@@ -14,16 +14,24 @@ module Braintrust
       return metrics unless usage_hash.is_a?(Hash)
 
       # Field mappings: OpenAI → Braintrust
+      # Supports both Chat Completions API (prompt_tokens, completion_tokens)
+      # and Responses API (input_tokens, output_tokens)
       field_map = {
         "prompt_tokens" => "prompt_tokens",
         "completion_tokens" => "completion_tokens",
-        "total_tokens" => "tokens"
+        "total_tokens" => "tokens",
+        # Responses API uses different field names
+        "input_tokens" => "prompt_tokens",
+        "output_tokens" => "completion_tokens"
       }
 
       # Prefix mappings for *_tokens_details
       prefix_map = {
         "prompt" => "prompt",
-        "completion" => "completion"
+        "completion" => "completion",
+        # Responses API uses input/output prefixes
+        "input" => "prompt",
+        "output" => "completion"
       }
 
       usage_hash.each do |key, value|

data/lib/braintrust/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module Braintrust
-  VERSION = "0.0.7"
+  VERSION = "0.0.8"
 end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: braintrust
 version: !ruby/object:Gem::Version
-  version: 0.0.7
+  version: 0.0.8
 platform: ruby
 authors:
 - Braintrust