braintrust 0.1.3 → 0.1.4

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 83ff9b69dc144333dba85a5f68e5a40d482ca99b3cae4bb55abe24ef2d05c296
-  data.tar.gz: 1a52913de27b3536c7881203f91d3d6050d53e66afeb90900d3c8a04b180951d
+  metadata.gz: e1a5c8840f707c7b4da95e4ccc8abea32591606d667a309432f2955d5df26eca
+  data.tar.gz: a45e62f34a1d59dd11e1cc46ff8d128a495a45a80fa1ce2026c76b648b58de89
 SHA512:
-  metadata.gz: fb7da28ba278c6a1cff5bd143e28808c723f1bf1507a6fe73d55b76f81d17e74ffc62f5c5dde030a1a5101797f3399592d13d525d46ad39b6b96047f7e47a3d6
-  data.tar.gz: d49db21d70faba9e3e9b59a61b88d55897cd5420ef636ae18c7b68a4c50d1428be2fbe0c0efc0b26f3a159bdbfb629025f0cc1aa8489187ecf5b586d57c8e1d2
+  metadata.gz: 75b71465a80ed2cfd3c6600113dd62357d01e0bd672f2043045f56d7d0223882cc2c5fd9f8927973ae546f99b68971bbe66fd34d66ec6fd62fafd65ca52abcd7
+  data.tar.gz: 06eb21fec07c05755aacd0a214cd16594f47a7033e723a2798df26692a0c15ccd9d5ed614588cc805620d303bb0f1a1c577c66b25c92c6dbbaa40283023cd662

data/lib/braintrust/api/datasets.rb

@@ -4,6 +4,7 @@ require "net/http"
 require "json"
 require "uri"
 require_relative "../logger"
+require_relative "../internal/http"
 
 module Braintrust
   class API
@@ -111,6 +112,7 @@ module Braintrust
         payload[:version] = version if version
 
         response = http_post_json_raw("/btql", payload)
+        Braintrust::Internal::Http.decompress_response!(response)
 
         # Parse JSONL response
         records = response.body.lines
@@ -158,9 +160,7 @@ module Braintrust
         start_time = Time.now
         Log.debug("[API] #{method.upcase} #{uri}")
 
-        http = Net::HTTP.new(uri.host, uri.port)
-        http.use_ssl = (uri.scheme == "https")
-        response = http.request(request)
+        response = Braintrust::Internal::Http.with_redirects(uri, request)
 
         duration_ms = ((Time.now - start_time) * 1000).round(2)
         Log.debug("[API] #{method.upcase} #{uri} -> #{response.code} (#{duration_ms}ms, #{response.body.bytesize} bytes)")

data/lib/braintrust/api/functions.rb

@@ -4,6 +4,7 @@ require "net/http"
 require "json"
 require "uri"
 require_relative "../logger"
+require_relative "../internal/http"
 
 module Braintrust
   class API
@@ -242,9 +243,7 @@ module Braintrust
         start_time = Time.now
         Log.debug("[API] #{method.upcase} #{uri}")
 
-        http = Net::HTTP.new(uri.host, uri.port)
-        http.use_ssl = (uri.scheme == "https")
-        response = http.request(request)
+        response = Braintrust::Internal::Http.with_redirects(uri, request)
 
         duration_ms = ((Time.now - start_time) * 1000).round(2)
         Log.debug("[API] #{method.upcase} #{uri} -> #{response.code} (#{duration_ms}ms, #{response.body.bytesize} bytes)")

data/lib/braintrust/api/internal/auth.rb

@@ -4,6 +4,7 @@ require "net/http"
 require "json"
 require "uri"
 require_relative "../../logger"
+require_relative "../../internal/http"
 
 module Braintrust
   class API
@@ -44,12 +45,7 @@ module Braintrust
           request = Net::HTTP::Post.new(uri)
           request["Authorization"] = "Bearer #{api_key}"
 
-          http = Net::HTTP.new(uri.hostname, uri.port)
-          http.use_ssl = true if uri.scheme == "https"
-
-          response = http.start do |http_session|
-            http_session.request(request)
-          end
+          response = Braintrust::Internal::Http.with_redirects(uri, request)
 
           Log.debug("Login: received response [#{response.code}]")
 

data/lib/braintrust/api/internal/experiments.rb

@@ -3,6 +3,7 @@
 require "net/http"
 require "json"
 require "uri"
+require_relative "../../internal/http"
 
 module Braintrust
   class API
@@ -22,7 +23,8 @@ module Braintrust
         # @param tags [Array<String>, nil] Optional tags
         # @param metadata [Hash, nil] Optional metadata
         # @return [Hash] Experiment data with "id", "name", "project_id", etc.
-        def create(name:, project_id:, ensure_new: true, tags: nil, metadata: nil)
+        def create(name:, project_id:, ensure_new: true, tags: nil, metadata: nil,
+          dataset_id: nil, dataset_version: nil)
           uri = URI("#{@state.api_url}/v1/experiment")
 
           payload = {
@@ -32,15 +34,15 @@ module Braintrust
           }
           payload[:tags] = tags if tags
           payload[:metadata] = metadata if metadata
+          payload[:dataset_id] = dataset_id if dataset_id
+          payload[:dataset_version] = dataset_version if dataset_version
 
           request = Net::HTTP::Post.new(uri)
           request["Content-Type"] = "application/json"
           request["Authorization"] = "Bearer #{@state.api_key}"
           request.body = JSON.dump(payload)
 
-          http = Net::HTTP.new(uri.host, uri.port)
-          http.use_ssl = (uri.scheme == "https")
-          response = http.request(request)
+          response = Braintrust::Internal::Http.with_redirects(uri, request)
 
           unless response.is_a?(Net::HTTPSuccess)
             raise Error, "HTTP #{response.code} for POST #{uri}: #{response.body}"

data/lib/braintrust/api/internal/projects.rb

@@ -3,6 +3,7 @@
 require "net/http"
 require "json"
 require "uri"
+require_relative "../../internal/http"
 
 module Braintrust
   class API
@@ -26,9 +27,7 @@ module Braintrust
           request["Authorization"] = "Bearer #{@state.api_key}"
           request.body = JSON.dump({name: name})
 
-          http = Net::HTTP.new(uri.host, uri.port)
-          http.use_ssl = (uri.scheme == "https")
-          response = http.request(request)
+          response = Braintrust::Internal::Http.with_redirects(uri, request)
 
           unless response.is_a?(Net::HTTPSuccess)
             raise Error, "HTTP #{response.code} for POST #{uri}: #{response.body}"

data/lib/braintrust/eval/functions.rb

@@ -98,9 +98,18 @@ module Braintrust
                 # The remote scorer receives all scorer arguments
                 result = api.functions.invoke(id: function_id, input: scorer_input)
 
-                # Parse result as float score
-                # The remote function should return a number
-                score = result.is_a?(Numeric) ? result.to_f : result.to_s.to_f
+                score = case result
+                when Hash
+                  if result.key?("score")
+                    result["score"].to_f
+                  else
+                    raise Error, "Hash result must contain 'score' key"
+                  end
+                when String
+                  result.to_f
+                else
+                  raise Error, "Unsupported result type: #{result.class}"
+                end
 
                 span.set_attribute("braintrust.output_json", JSON.dump(score))
                 score

data/lib/braintrust/eval.rb

@@ -220,8 +220,14 @@ module Braintrust
         api.login
 
         # Resolve dataset to cases if dataset parameter provided
+        dataset_id = nil
+        dataset_version = nil
+
         if dataset
-          cases = resolve_dataset(dataset, project, api)
+          resolved = resolve_dataset(dataset, project, api)
+          cases = resolved[:cases]
+          dataset_id = resolved[:dataset_id]
+          dataset_version = resolved[:dataset_version]
         end
 
         # Register project and experiment via internal API
@@ -234,7 +240,9 @@ module Braintrust
           project_id: project_result["id"],
           ensure_new: !update,
           tags: tags,
-          metadata: metadata
+          metadata: metadata,
+          dataset_id: dataset_id,
+          dataset_version: dataset_version
         )
 
         experiment_id = experiment_result["id"]
@@ -292,11 +300,11 @@ module Braintrust
         end
       end
 
-      # Resolve dataset parameter to an array of case records
+      # Resolve dataset parameter to cases with metadata for experiment linking
       # @param dataset [String, Hash, Dataset] Dataset specifier or instance
       # @param project [String] Project name (used as default if not specified)
       # @param api [API] Braintrust API client
-      # @return [Array<Hash>] Array of case records
+      # @return [Hash] Hash with :cases, :dataset_id, and :dataset_version
       def resolve_dataset(dataset, project, api)
         limit = nil
 
@@ -315,7 +323,15 @@ module Braintrust
           raise ArgumentError, "dataset must be String, Hash, or Dataset, got #{dataset.class}"
         end
 
-        dataset_obj.fetch_all(limit: limit)
+        cases = dataset_obj.fetch_all(limit: limit)
+
+        # Use pinned version if available, otherwise compute from max(_xact_id)
+        version = dataset_obj.version
+        version ||= cases
+          .filter_map { |c| c[:origin] && JSON.parse(c[:origin])["_xact_id"] }
+          .max
+
+        {cases: cases, dataset_id: dataset_obj.id, dataset_version: version}
       end
     end
   end

data/lib/braintrust/internal/http.rb

@@ -0,0 +1,97 @@
+# frozen_string_literal: true
+
+require "net/http"
+require "uri"
+require "zlib"
+require "stringio"
+require_relative "../logger"
+
+module Braintrust
+  module Internal
+    # HTTP utilities for redirect following and response decompression.
+    # Drop-in enhancement for raw Net::HTTP request calls throughout the SDK.
+    module Http
+      DEFAULT_MAX_REDIRECTS = 5
+
+      # Execute an HTTP request, following redirects as needed.
+      #
+      # @param uri [URI] The request URI
+      # @param request [Net::HTTPRequest] The prepared request object
+      # @param max_redirects [Integer] Maximum number of redirects to follow
+      # @return [Net::HTTPResponse] The final response
+      # @raise [Braintrust::Error] On too many redirects or missing Location header
+      def self.with_redirects(uri, request, max_redirects: DEFAULT_MAX_REDIRECTS)
+        response = perform_request(uri, request)
+
+        redirects = 0
+        original_request = request
+
+        while response.is_a?(Net::HTTPRedirection)
+          redirects += 1
+          if redirects > max_redirects
+            raise Error, "Too many redirects (max #{max_redirects})"
+          end
+
+          location = response["location"]
+          unless location
+            raise Error, "Redirect response #{response.code} without Location header"
+          end
+
+          redirect_uri = URI(location)
+          redirect_uri = uri + redirect_uri unless redirect_uri.host
+
+          Log.debug("[HTTP] Following #{response.code} redirect to #{redirect_uri}")
+
+          request = build_redirect_request(response, redirect_uri, original_request, uri)
+          uri = redirect_uri
+          response = perform_request(uri, request)
+        end
+
+        response
+      end
+
+      # Decompress an HTTP response body in place based on Content-Encoding.
+      # No-op if the response has no recognized encoding.
+      #
+      # @param response [Net::HTTPResponse] The response to decompress
+      # @return [void]
+      def self.decompress_response!(response)
+        encoding = response["content-encoding"]&.downcase
+        case encoding
+        when "gzip", "x-gzip"
+          gz = Zlib::GzipReader.new(StringIO.new(response.body))
+          response.body.replace(gz.read)
+          gz.close
+          response.delete("content-encoding")
+        end
+      end
+
+      def self.perform_request(uri, request)
+        http = Net::HTTP.new(uri.host, uri.port)
+        http.use_ssl = (uri.scheme == "https")
+        http.request(request)
+      end
+      private_class_method :perform_request
+
+      def self.build_redirect_request(response, redirect_uri, original_request, original_uri)
+        if response.code == "307" || response.code == "308"
+          request = original_request.class.new(redirect_uri)
+          request.body = original_request.body
+          request["Content-Type"] = original_request["Content-Type"] if original_request["Content-Type"]
+        else
+          # 301, 302, 303: follow with GET, no body
+          request = Net::HTTP::Get.new(redirect_uri)
+        end
+
+        # Strip Authorization when redirecting to a different host (e.g. S3)
+        if original_uri.host == redirect_uri.host
+          auth = original_request["Authorization"]
+          request["Authorization"] = auth if auth
+        end
+
+        request
+      end
+      private_class_method :build_redirect_request
+    end
+  end
+end

data/lib/braintrust/trace/attachment.rb

@@ -2,6 +2,7 @@
 
 require "net/http"
 require_relative "../internal/encoding"
+require_relative "../internal/http"
 require "uri"
 
 module Braintrust
@@ -91,7 +92,8 @@ module Braintrust
       #   att = Braintrust::Trace::Attachment.from_url("https://example.com/image.png")
       def self.from_url(url)
         uri = URI.parse(url)
-        response = Net::HTTP.get_response(uri)
+        request = Net::HTTP::Get.new(uri)
+        response = Braintrust::Internal::Http.with_redirects(uri, request)
 
         unless response.is_a?(Net::HTTPSuccess)
           raise StandardError, "Failed to fetch URL: #{response.code} #{response.message}"

data/lib/braintrust/version.rb

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

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: braintrust
 version: !ruby/object:Gem::Version
-  version: 0.1.3
+  version: 0.1.4
 platform: ruby
 authors:
 - Braintrust
@@ -242,6 +242,7 @@ files:
 - lib/braintrust/eval/summary.rb
 - lib/braintrust/internal/encoding.rb
 - lib/braintrust/internal/env.rb
+- lib/braintrust/internal/http.rb
 - lib/braintrust/internal/origin.rb
 - lib/braintrust/internal/template.rb
 - lib/braintrust/internal/thread_pool.rb