bitfab 0.38.3 → 0.40.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 28a82b98d8327aead543061c18b95a22b2100f89e0289ffdd6ad85a262e08447
-  data.tar.gz: 18c101e18d509cc2a5fb880210d3d5cfaee0f9311c7521abeaa2716bdeca563f
+  metadata.gz: a178d05e7c40055c3ac8dec21a40eb60eb36b08a005dfa55d031c3c3b9cdf129
+  data.tar.gz: e71a921a9bad32b967c723245c1187257366073f13aab6f8dca2882d6b758fb8
 SHA512:
-  metadata.gz: 960e4f390e26e953cb555a1a359fd7e3b505172a74e513fe69f281dc7e1111b03e8ac93451154d632dd225e36cdf33f5c54dc3dc59dbd7379a7af3672f7fb87b
-  data.tar.gz: ac599caf319a0dc2814147faee28efd5940609bdc2200d696db4bcb3581fc72fafe443bb50f65204ffb46eea5fa8f0ff9806dd78c42a8d3599e1429263d2ab71
+  metadata.gz: 559fed18dd194a527ff477e057f465a172d7266ba05c4e6b3305a693091b0fa1de3847078d8f0df4c6822a4b28b2fd6a0f802eaf2db7d75e38233051f05fc887
+  data.tar.gz: f4a718c47568294809eaed27905b2d8b67a3058fc954eff2dac6b10a05ae2219463d1742b40bb0a7725d4979a5c86c2beb37e9e10621a588e007a9a7757d6376

data/lib/bitfab/client.rb

@@ -13,7 +13,7 @@ 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
@@ -30,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
@@ -43,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.
@@ -157,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
@@ -203,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]
@@ -223,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
@@ -232,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:}
@@ -281,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
@@ -304,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
@@ -357,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
@@ -394,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
@@ -410,6 +436,7 @@ 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
@@ -551,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
@@ -690,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
@@ -772,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}"
@@ -819,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
 
@@ -826,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:})
@@ -847,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"]
@@ -955,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/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module Bitfab
-  VERSION = "0.38.3"
+  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.3
+  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