bitfab 0.38.2 → 0.40.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 29901e7611834921e9378870830af40c08e5af02d1ca0708585a02b169b12e6e
-  data.tar.gz: b42984a3710e9399159fd44af3aaac7b76e3155bd3121eaa5482aa27570fbb2e
+  metadata.gz: a178d05e7c40055c3ac8dec21a40eb60eb36b08a005dfa55d031c3c3b9cdf129
+  data.tar.gz: e71a921a9bad32b967c723245c1187257366073f13aab6f8dca2882d6b758fb8
 SHA512:
-  metadata.gz: 576a44cbe7e2fdf7270f14af3d8a98541bb2e15cd2b02289c18d58f033688d6797b2b3a0cfc8153f1ba8adbf0d72f9b9b9241e1bfd5fbd3ffcc687a4cb9adb44
-  data.tar.gz: ce3586fbdc541ffad387d87ca4299ee34df5e0c61a5c01a6b5112ee11084cb7900d6616f9e3db45ec0386fc019bc1712f267ea3c5799809880ae3611d3d30cb9
+  metadata.gz: 559fed18dd194a527ff477e057f465a172d7266ba05c4e6b3305a693091b0fa1de3847078d8f0df4c6822a4b28b2fd6a0f802eaf2db7d75e38233051f05fc887
+  data.tar.gz: f4a718c47568294809eaed27905b2d8b67a3058fc954eff2dac6b10a05ae2219463d1742b40bb0a7725d4979a5c86c2beb37e9e10621a588e007a9a7757d6376

data/lib/bitfab/client.rb

@@ -8,11 +8,12 @@ require_relative "http_client"
 require_relative "replay"
 require_relative "span_context"
 require_relative "serialize"
+require_relative "timestamp"
 require_relative "warn_once"
 
 module Bitfab
   class Client
-    SPAN_TYPES = %w[llm agent function guardrail handoff custom].freeze
+    SPAN_TYPES = ["llm", "agent", "function", "guardrail", "handoff", "custom"].freeze
     UUID_PATTERN = /\A[0-9a-f]{8}-[0-9a-f]{4}-[1-8][0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}\z/i
 
     # Sentinel returned by check_mock_replay when this span should run real
@@ -29,6 +30,9 @@ module Bitfab
 
     attr_reader :service_url
 
+    # Dataset operations for the authenticated organization.
+    attr_reader :datasets
+
     def initialize(api_key: nil, service_url: nil, enabled: true, strict: false)
       @api_key_config = api_key
       @service_url = service_url || DEFAULT_SERVICE_URL
@@ -42,6 +46,7 @@ module Bitfab
       # The key is NOT read here. HttpClient gets a proc so the key is resolved
       # at send time, after any env loading has run.
       @http_client = HttpClient.new(api_key: -> { resolve_api_key }, service_url: @service_url)
+      @datasets = Datasets.new(@http_client)
       # Mock overrides registered via register_mock_override, applied to every
       # replay on this client (after any per-call mock_override). Instance
       # state, no global; clear_mock_overrides resets it.
@@ -156,10 +161,25 @@ module Bitfab
       name: nil, code_change_description: Replay::CODE_CHANGE_UNSET, code_change_files: Replay::CODE_CHANGE_UNSET, experiment_group_id: nil, dataset_id: nil, grader_ids: nil, mock: "marked",
       adapt_inputs: nil, mock_override: nil, db_branch: nil, on_item_start: nil, on_item_finish: nil, on_progress: nil)
       Replay.run(
-        self, receiver, method_name,
-        trace_function_key:, limit:, trace_ids:, name:, max_concurrency:,
-        code_change_description:, code_change_files:, experiment_group_id:, dataset_id:, grader_ids:, mock:, adapt_inputs:,
-        mock_override:, db_branch:, on_item_start:, on_item_finish:,
+        self,
+        receiver,
+        method_name,
+        trace_function_key:,
+        limit:,
+        trace_ids:,
+        name:,
+        max_concurrency:,
+        code_change_description:,
+        code_change_files:,
+        experiment_group_id:,
+        dataset_id:,
+        grader_ids:,
+        mock:,
+        adapt_inputs:,
+        mock_override:,
+        db_branch:,
+        on_item_start:,
+        on_item_finish:,
         on_progress:
       )
     end
@@ -202,18 +222,18 @@ module Bitfab
               match: ->(node) { node[:trace_function_key] == trace_function_key && keyed_match.call(node) },
               value: keyed_override[:value]
             }
-            return nil
+            return
           end
           if keyed_override.respond_to?(:call)
             @mock_overrides << {
               match: ->(node) { node[:trace_function_key] == trace_function_key },
               value: keyed_override
             }
-            return nil
+            return
           end
           raise ArgumentError,
             "register_mock_override(trace_function_key, override) requires " \
-            "a { match:, value: } hash or callable resolver."
+              "a { match:, value: } hash or callable resolver."
         elsif positional.length == 1 && positional[0].is_a?(Hash)
           override = positional[0]
           match = override[:match]
@@ -222,7 +242,7 @@ module Bitfab
           value = override.key?(:value) ? override[:value] : VALUE_UNSET
         elsif positional.length == 1 && positional[0].respond_to?(:call)
           @mock_overrides << {match: ->(_) { true }, value: positional[0]}
-          return nil
+          return
         elsif positional.length == 2
           match, value = positional
         end
@@ -231,16 +251,16 @@ module Bitfab
       unless match.respond_to?(:call)
         raise ArgumentError,
           "register_mock_override requires a callable match. Pass (match, value) " \
-          "positionally, as keywords (match:, value:), or as a { match:, value: } hash. " \
-          "value may be a flat value or a callable."
+            "positionally, as keywords (match:, value:), or as a { match:, value: } hash. " \
+            "value may be a flat value or a callable."
       end
       # A forgotten value must not silently inject nil. An explicit nil is a
       # legitimate injected value and passes this guard.
       if value.equal?(VALUE_UNSET)
         raise ArgumentError,
           "register_mock_override requires a value (the second argument, or " \
-          "value:). It may be a flat value or a callable; pass value: nil " \
-          "explicitly to inject nil."
+            "value:). It may be a flat value or a callable; pass value: nil " \
+            "explicitly to inject nil."
       end
 
       @mock_overrides << {match:, value:}
@@ -280,12 +300,13 @@ module Bitfab
     def get_trace_span(trace_id, id: nil, name: nil, occurrence: "last")
       validate_trace_id(trace_id)
       raise ArgumentError, "Provide exactly one of id or name" if id.nil? == name.nil?
+
       validate_span_id(id) unless id.nil?
       if !name.nil? && (!name.is_a?(String) || name.empty?)
         raise ArgumentError, "name must be a non-empty string"
       end
 
-      valid_occurrence = %w[first last].include?(occurrence) || (occurrence.is_a?(Integer) && occurrence >= 0)
+      valid_occurrence = ["first", "last"].include?(occurrence) || (occurrence.is_a?(Integer) && occurrence >= 0)
       unless valid_occurrence
         raise ArgumentError, 'occurrence must be "first", "last", or a non-negative integer'
       end
@@ -303,11 +324,11 @@ module Bitfab
       return yield unless tracing_enabled?
 
       resolved_capture_when = capture_when.to_s
-      unless %w[always nested].include?(resolved_capture_when)
+      unless ["always", "nested"].include?(resolved_capture_when)
         Bitfab.warn_once(
           "invalid-capture-when:#{trace_function_key}",
           "unknown capture_when value #{capture_when.inspect}; defaulting to \"always\". " \
-          "Valid values: \"always\", \"nested\"."
+            "Valid values: \"always\", \"nested\"."
         )
         resolved_capture_when = "always"
       end
@@ -333,7 +354,7 @@ module Bitfab
         span_id = SecureRandom.uuid
         parent_span_id = parent&.dig(:span_id)
         is_root_span = parent_span_id.nil?
-        started_at = Time.now.utc.strftime("%Y-%m-%dT%H:%M:%S.%3NZ")
+        started_at = Bitfab.now_iso_timestamp
         resolved_test_run_id = replay_ctx&.dig(:test_run_id)
         resolved_input_source_span_id = replay_ctx&.dig(:input_source_span_id)
         resolved_input_source_trace_id = replay_ctx&.dig(:input_source_trace_id)
@@ -356,8 +377,14 @@ module Bitfab
         call_index = advance_mock_counter(replay_ctx, trace_function_key, span_name, is_root_span:)
         if call_index
           mock_result = check_mock_replay(
-            replay_ctx, trace_function_key, span_name, call_index,
-            span_type:, args:, kwargs:, mock_on_replay:
+            replay_ctx,
+            trace_function_key,
+            span_name,
+            call_index,
+            span_type:,
+            args:,
+            kwargs:,
+            mock_on_replay:
           )
           if mock_result != MOCK_REPLAY_MISS
             mocked_output, mock_source = mock_result
@@ -393,7 +420,7 @@ module Bitfab
         Bitfab.warn_once(
           "span-setup:#{trace_function_key}",
           "span setup failed for '#{trace_function_key}'; this call runs untraced. " \
-          "Your method still executes and returns normally."
+            "Your method still executes and returns normally."
         )
         return yield
       end
@@ -409,10 +436,11 @@ module Bitfab
         # only the first call sends the span. Subsequent calls (e.g. from the
         # enumerator wrapper after iteration completes) are no-ops.
         next if finalized
+
         finalized = true
 
         begin
-          ended_at = Time.now.utc.strftime("%Y-%m-%dT%H:%M:%S.%3NZ")
+          ended_at = Bitfab.now_iso_timestamp
 
           send_span(
             trace_function_key:,
@@ -550,7 +578,7 @@ module Bitfab
       end
       if @explicitly_enabled && !@api_key_warned
         @api_key_warned = true
-        warn "Bitfab: api_key is empty: tracing is disabled. Provide a valid API key to enable tracing."
+        warn("Bitfab: api_key is empty: tracing is disabled. Provide a valid API key to enable tracing.")
       end
       nil
     end
@@ -689,7 +717,7 @@ module Bitfab
       # still rides out with dropped: true, so the server scrubs any sibling
       # spans that already raced out before the flag was set.
       trace_dropped = TraceState.get(trace_id)&.dig(:dropped) == true
-      return nil if trace_dropped
+      return if trace_dropped
 
       # Human-readable JSON (input/output fields). The *_with_report variants
       # also report what could not be faithfully captured so a lossy span is
@@ -771,8 +799,8 @@ module Bitfab
     # mocked) so unmarked spans don't silently shift subsequent marked
     # spans' indices. Different (key, name) pairs have independent counters.
     def advance_mock_counter(replay_ctx, trace_function_key, span_name, is_root_span:)
-      return nil if is_root_span
-      return nil unless replay_ctx&.dig(:mock_tree)
+      return if is_root_span
+      return unless replay_ctx&.dig(:mock_tree)
 
       counters = replay_ctx[:call_counters]
       counter_key = "#{trace_function_key}:#{span_name}"
@@ -818,6 +846,7 @@ module Bitfab
           # output, wrap it in a callable value.
           unless value.respond_to?(:call)
             return [value, "override"] unless value.equal?(NO_MOCK_OVERRIDE)
+
             next
           end
 
@@ -825,6 +854,7 @@ module Bitfab
             unless mock_entry
               raise "No recorded span to source output for '#{trace_function_key}'."
             end
+
             resolve_recorded_output(mock_entry, replay_ctx)
           end
           resolved = value.call({node:, inputs: args, kwargs:, get_original_output:})
@@ -846,7 +876,7 @@ module Bitfab
       unless mock_entry
         raise StandardError,
           "Replay selected span '#{trace_function_key}:#{span_name}' for mocking, " \
-          "but recorded occurrence #{call_index + 1} is unavailable. The real span was not executed."
+            "but recorded occurrence #{call_index + 1} is unavailable. The real span was not executed."
       end
 
       [resolve_recorded_output(mock_entry, replay_ctx), "recorded"]
@@ -889,7 +919,7 @@ module Bitfab
     def send_mocked_span(trace_function_key:, trace_id:, span_id:, parent_span_id:,
       span_name:, span_type:, function_name:, args:, kwargs:, mocked_output:,
       started_at:, test_run_id:, input_source_span_id:, mock_source:)
-      ended_at = Time.now.utc.strftime("%Y-%m-%dT%H:%M:%S.%3NZ")
+      ended_at = Bitfab.now_iso_timestamp
       send_span(
         trace_function_key:,
         trace_id:,
@@ -954,9 +984,13 @@ module Bitfab
     #   mock strategy. A missing selected occurrence fails closed.
     def wrap(klass, method_name, name: nil, type: "custom", capture_when: "always", mock_on_replay: false)
       Bitfab::Traceable.wrap(
-        klass, method_name,
+        klass,
+        method_name,
         trace_function_key: @trace_function_key,
-        name:, type:, capture_when:, mock_on_replay:,
+        name:,
+        type:,
+        capture_when:,
+        mock_on_replay:,
         client: @client
       )
     end

data/lib/bitfab/datasets.rb

@@ -0,0 +1,107 @@
+# frozen_string_literal: true
+
+module Bitfab
+  # Dataset operations for the authenticated organization, reached as
+  # `client.datasets`. A dataset is a named bucket of traces scoped to one
+  # trace function. Experiments replay against it and its graders score its
+  # members. Every method returns the parsed JSON response as a Hash with
+  # string keys matching the HTTP API ("traceCount", "addedTraceIds", ...).
+  class Datasets
+    DEFAULT_RERUN_TIMEOUT_SECONDS = 90
+    DEFAULT_RERUN_POLL_INTERVAL_SECONDS = 1
+    TERMINAL_RERUN_STATUSES = ["completed", "errored"].freeze
+
+    def initialize(http_client)
+      @http_client = http_client
+    end
+
+    # Create a dataset, or update the one already named this way under the
+    # same trace function. The result's "created" reports which happened. A
+    # nil description leaves an existing description untouched.
+    def save(trace_function_key:, name:, description: nil)
+      payload = {"traceFunctionKey" => trace_function_key, "name" => name}
+      payload["description"] = description unless description.nil?
+      @http_client.request("/api/sdk/datasets", payload)
+    end
+
+    # List datasets, scoped to one trace function when given and
+    # organization-wide otherwise.
+    def list(trace_function_key: nil)
+      query = trace_function_key.nil? ? "" : "?#{URI.encode_www_form("traceFunctionKey" => trace_function_key)}"
+      @http_client.get("/api/sdk/datasets#{query}")["datasets"]
+    end
+
+    # Fetch one dataset by id. A dataset outside this organization raises
+    # Net::HTTPError with a 404.
+    def get(dataset_id)
+      @http_client.get(dataset_path(dataset_id))["dataset"]
+    end
+
+    # The ids of every trace in the dataset, the same membership a replay with
+    # dataset_id: selects.
+    def list_traces(dataset_id)
+      @http_client.get(dataset_path(dataset_id, "/traces"))
+    end
+
+    # Add traces to the dataset (1 to 100 ids per call). Traces outside the
+    # organization or under another trace function are reported in
+    # "skippedTraceIds" rather than failing the call.
+    def add_traces(dataset_id, trace_ids)
+      @http_client.request(dataset_path(dataset_id, "/traces"), {"traceIds" => trace_ids})
+    end
+
+    # Remove traces from the dataset. The traces themselves are never deleted.
+    def remove_traces(dataset_id, trace_ids)
+      @http_client.request(dataset_path(dataset_id, "/removeTraces"), {"traceIds" => trace_ids})
+    end
+
+    # Assign graders to the dataset (1 to 100 ids per call). Graders outside
+    # the organization or under another trace function are reported in
+    # "skippedGraderIds" rather than failing the call.
+    def add_graders(dataset_id, grader_ids)
+      @http_client.request(dataset_path(dataset_id, "/graders"), {"graderIds" => grader_ids})
+    end
+
+    # Unassign graders from the dataset.
+    def remove_graders(dataset_id, grader_ids)
+      @http_client.request(dataset_path(dataset_id, "/removeGraders"), {"graderIds" => grader_ids})
+    end
+
+    # Re-run graders over every trace in the dataset. Defaults to every
+    # assigned grader, and an unassigned id is rejected. Waits for the run to
+    # finish (up to `timeout` seconds) unless `wait:` is false, and returns the
+    # last run state seen either way. A request matching an in-flight run
+    # joins it.
+    def rerun_graders(dataset_id, grader_ids: nil, wait: true,
+      timeout: DEFAULT_RERUN_TIMEOUT_SECONDS, poll_interval: DEFAULT_RERUN_POLL_INTERVAL_SECONDS)
+      payload = grader_ids.nil? ? {} : {"graderIds" => grader_ids}
+      started = @http_client.request(dataset_path(dataset_id, "/rerunGraders"), payload)
+      return started unless wait
+
+      deadline = monotonic_now + timeout
+      run = started["run"]
+      while !TERMINAL_RERUN_STATUSES.include?(run["status"]) && monotonic_now < deadline
+        sleep(poll_interval)
+        run = get_grader_rerun(dataset_id, run_id: run["id"]) || run
+      end
+      {"run" => run, "joinedExisting" => started["joinedExisting"]}
+    end
+
+    # The dataset's active grader re-run, or the run named by run_id:. Returns
+    # nil when nothing is active or the run is not this dataset's.
+    def get_grader_rerun(dataset_id, run_id: nil)
+      query = run_id.nil? ? "" : "?#{URI.encode_www_form("runId" => run_id)}"
+      @http_client.get(dataset_path(dataset_id, "/rerunGraders#{query}"))["run"]
+    end
+
+    private
+
+    def dataset_path(dataset_id, suffix = "")
+      "/api/sdk/datasets/#{URI.encode_www_form_component(dataset_id)}#{suffix}"
+    end
+
+    def monotonic_now
+      Process.clock_gettime(Process::CLOCK_MONOTONIC)
+    end
+  end
+end

data/lib/bitfab/span_context.rb

@@ -1,5 +1,7 @@
 # frozen_string_literal: true
 
+require_relative "timestamp"
+
 module Bitfab
   # Handle to the current active span, allowing context to be added.
   class CurrentSpan
@@ -171,7 +173,7 @@ module Bitfab
     def create(trace_id, test_run_id: nil, input_source_trace_id: nil)
       @states_mutex.synchronize do
         @states[trace_id] ||= begin
-          started_at = Time.now.utc.strftime("%Y-%m-%dT%H:%M:%S.%3NZ")
+          started_at = Bitfab.now_iso_timestamp
           {
             trace_id:,
             started_at:,

data/lib/bitfab/timestamp.rb

@@ -0,0 +1,16 @@
+# frozen_string_literal: true
+
+require "time"
+
+module Bitfab
+  @timestamp_mutex = Mutex.new
+  @last_timestamp_us = 0
+
+  def self.now_iso_timestamp
+    wall_clock_us = (Time.now.to_r * 1_000_000).to_i
+    timestamp_us = @timestamp_mutex.synchronize do
+      @last_timestamp_us = [wall_clock_us, @last_timestamp_us + 1].max
+    end
+    Time.at(Rational(timestamp_us, 1_000_000)).utc.strftime("%Y-%m-%dT%H:%M:%S.%6NZ")
+  end
+end

data/lib/bitfab/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module Bitfab
-  VERSION = "0.38.2"
+  VERSION = "0.40.0"
 end

data/lib/bitfab.rb

@@ -10,6 +10,7 @@ require_relative "bitfab/serialize"
 require_relative "bitfab/db_snapshot"
 require_relative "bitfab/span_context"
 require_relative "bitfab/http_client"
+require_relative "bitfab/datasets"
 require_relative "bitfab/mock_override"
 require_relative "bitfab/replay"
 require_relative "bitfab/replay_branch"
@@ -119,16 +120,16 @@ module Bitfab
     #   whose source trace carried no DB snapshot reference
     def current_replay_branch
       ctx = ReplayContext.current
-      return nil unless ctx
+      return unless ctx
 
       lease = ctx[:db_branch_lease]
-      return nil unless lease
+      return unless lease
 
       # Surface the Bitfab trace ID (what the customer sees in the dashboard),
       # not the external one. Falling back to the external ID keeps replays from
       # external sources working until that path is fully wired.
       trace_id = ctx[:source_bitfab_trace_id] || ctx[:input_source_trace_id]
-      return nil unless trace_id
+      return unless trace_id
 
       ReplayBranch.new(lease, trace_id, ctx)
     end
@@ -159,7 +160,7 @@ module Bitfab
     #   succeeded, errored, and item (the object replay already passes to
     #   on_item_finish)
     def report_replay_progress(progress)
-      warn "#{BITFAB_PROGRESS_PREFIX}#{Replay.json_safe(progress).to_json}"
+      warn("#{BITFAB_PROGRESS_PREFIX}#{Replay.json_safe(progress).to_json}")
     rescue
       nil
     end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: bitfab
 version: !ruby/object:Gem::Version
-  version: 0.38.2
+  version: 0.40.0
 platform: ruby
 authors:
 - Harvest Team
@@ -142,6 +142,7 @@ files:
 - lib/bitfab/client.rb
 - lib/bitfab/compress.rb
 - lib/bitfab/constants.rb
+- lib/bitfab/datasets.rb
 - lib/bitfab/db_snapshot.rb
 - lib/bitfab/http_client.rb
 - lib/bitfab/mock_override.rb
@@ -153,6 +154,7 @@ files:
 - lib/bitfab/replay_registry.rb
 - lib/bitfab/serialize.rb
 - lib/bitfab/span_context.rb
+- lib/bitfab/timestamp.rb
 - lib/bitfab/traceable.rb
 - lib/bitfab/transport.rb
 - lib/bitfab/version.rb