bitfab 0.66.0 → 0.67.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 35b3530a206cb86806f06bd4cbeef8d660b07f118f656dae1bc81d9bb5418038
-  data.tar.gz: 3da107dd35bde42d9cd9f4c7f549226cf52624da941a24333c6b2ff7738dfd61
+  metadata.gz: 455ac9a2dfa11b39c17550b3f86f6ffa87f3ea0a9684fb032722184a705da7e0
+  data.tar.gz: 11c0f2491a21515f3fc75f4f4f8eb98f544b7e6d5d3c9d1c7fc09e0728787b16
 SHA512:
-  metadata.gz: 191854786921184a3324c8fc7aa5bb620c82310e28129e33986cde9aa77ddbdb9fdf0da917b7829a07ba501f2c3992b8d6e0742a7ca695abaa1e3dffcac69826
-  data.tar.gz: 9e642a61b8fb8287f6db422cf32d55a44565e9ff2caffa5716d0bf5800248780e620ba0f41ef1f10a35505079099f0f50fca519172185b559ba5c41256a99191
+  metadata.gz: d1f41b10072dd2b9b52f2a91ef5265a91805d2aa8129a05a30a303f6f55831934c8206eca6fd215ac938a526f499e9bc2b38a2e74476c8aa06e6c675e14d9b7f
+  data.tar.gz: d1e5d1ac8216e5505b5d5da628d6914349fe7ee6018fa55616b20a5ba3ca66e647484f4db176ee85feac28c476f702e831bfcf9ae3fa5ff5084d9d0570584ade

data/lib/bitfab/jev.rb

@@ -0,0 +1,149 @@
+# frozen_string_literal: true
+
+module Bitfab
+  module Jev
+    PROVIDER = "typesafe"
+    SYSTEM_ONE = :system_one
+    WRAPPER_SCOPE_KEY = :bitfab_jev_wrapper_scope
+
+    module_function
+
+    def jev_shaped?(reply)
+      return false if reply.nil? || reply.is_a?(String) || reply.is_a?(Numeric) || reply.is_a?(Symbol)
+
+      usage = field(reply, :usage)
+      field(reply, :model).is_a?(String) && !usage.nil? && !usage.is_a?(String) && !usage.is_a?(Numeric)
+    rescue
+      false
+    end
+
+    def auto_context(reply, duration_ms)
+      jev_shaped?(reply) ? context(reply, duration_ms) : nil
+    rescue
+      nil
+    end
+
+    def mark_recorded(reply)
+      Thread.current[WRAPPER_SCOPE_KEY]&.push(reply)
+    end
+
+    def watch_recorded
+      previous = Thread.current[WRAPPER_SCOPE_KEY]
+      scope = []
+      Thread.current[WRAPPER_SCOPE_KEY] = scope
+      reply = yield
+      [reply, !scope.empty?]
+    ensure
+      Thread.current[WRAPPER_SCOPE_KEY] = previous
+      previous&.concat(scope) if scope
+    end
+
+    def record_node_call(method_name)
+      return yield unless method_name.to_sym == SYSTEM_ONE
+
+      span = Bitfab.current_span
+      started_at = monotonic_ms
+      reply, recorded_inside = watch_recorded { yield }
+      begin
+        entry = recorded_inside ? nil : auto_context(reply, monotonic_ms - started_at)
+        if entry
+          span.add_context(entry)
+          mark_recorded(reply)
+          Subtree.active_sessions.each(&:jev_recorded_below!)
+        end
+      rescue
+        nil
+      end
+      reply
+    end
+
+    def field(value, name)
+      if value.is_a?(Hash)
+        value.key?(name) ? value[name] : value[name.to_s]
+      elsif value.respond_to?(name)
+        value.public_send(name)
+      end
+    end
+
+    def context(reply, duration_ms)
+      return nil if reply.nil?
+
+      usage = field(reply, :usage)
+      model = field(reply, :model)
+      input_tokens = usage.nil? ? nil : field(usage, :input_tokens)
+      output_tokens = usage.nil? ? nil : field(usage, :output_tokens)
+      entry = {provider: PROVIDER}
+      entry[:model] = model if model.is_a?(String)
+      entry[:inputTokens] = input_tokens if input_tokens.is_a?(Numeric)
+      entry[:outputTokens] = output_tokens if output_tokens.is_a?(Numeric)
+      entry[:durationMs] = duration_ms
+      entry
+    end
+
+    def record(span, reply, duration_ms)
+      entry = context(reply, duration_ms)
+      span.add_context(entry) if entry
+    rescue
+      nil
+    end
+
+    def monotonic_ms
+      Process.clock_gettime(Process::CLOCK_MONOTONIC, :millisecond)
+    end
+
+    class Wrapper
+      def initialize(callable, on_reply)
+        @callable = callable
+        @on_reply = on_reply
+      end
+
+      # Call the wrapped Jev method with the given arguments and return its
+      # reply unchanged. Errors from the call are raised as-is and record nothing.
+      def call(*args, **kwargs, &block)
+        span = Bitfab.current_span
+        started_at = Jev.monotonic_ms
+        reply, recorded = Jev.watch_recorded { @callable.call(*args, **kwargs, &block) }
+        Jev.record(span, reply, Jev.monotonic_ms - started_at) unless recorded
+        begin
+          @on_reply&.call(reply)
+        rescue
+          nil
+        end
+        reply
+      end
+
+      # A Proc that calls this wrapper, so it can be passed with `&`.
+      def to_proc
+        method(:call).to_proc
+      end
+    end
+  end
+
+  class Client
+    # Record the model, token usage, and duration of a Jev (TypeSafe System One)
+    # call on the span that is current when the call starts. Opens no span of
+    # its own and never raises from recording.
+    #
+    # @example
+    #   ask_jev = Bitfab.client.wrap_jev(jev.method(:system_one))
+    #   reply = ask_jev.call(state:, questions:)
+    #
+    # @param callable [#call, nil] calls Jev and returns its reply; pass a block instead to wrap the block
+    # @param on_reply [#call, nil] called with each reply after it is recorded; its errors are ignored
+    # @return [Bitfab::Jev::Wrapper] responds to +call+ with the same arguments as the wrapped callable
+    def wrap_jev(callable = nil, on_reply: nil, &block)
+      raise ArgumentError, "wrap_jev takes a callable or a block, not both." if callable && block
+      target = callable || block
+      raise ArgumentError, "wrap_jev requires a callable that calls Jev (e.g., jev.method(:system_one))." unless target.respond_to?(:call)
+
+      Jev::Wrapper.new(target, on_reply)
+    end
+  end
+
+  class BitfabFunction
+    # Jev wrappers record onto the currently active span under any key.
+    def wrap_jev(*args, **kwargs, &block)
+      @client.wrap_jev(*args, **kwargs, &block)
+    end
+  end
+end

data/lib/bitfab/subtree.rb

@@ -37,6 +37,7 @@ module Bitfab
       :name,
       :excluded,
       :passthrough,
+      :jev,
       keyword_init: true
     )
 
@@ -55,6 +56,8 @@ module Bitfab
       :declared_node,
       :contexts,
       :prompt,
+      :jev_started_at,
+      :jev_recorded_below,
       keyword_init: true
     )
 
@@ -506,7 +509,7 @@ module Bitfab
 
         return if event.event == :call && skip_boundary_call?(event)
 
-        resolved_path = owned_path(event.path)
+        resolved_path = traced_path(event)
         return unless resolved_path
 
         case event.event
@@ -523,6 +526,12 @@ module Bitfab
         @exclude[method_name.to_s] == true
       end
 
+      def jev_recorded_below!
+        @stacks[Fiber.current.object_id]&.each do |entry|
+          entry.jev_recorded_below = true if entry.is_a?(OpenSpan)
+        end
+      end
+
       def reserve_node
         stack = current_stack(Fiber.current.object_id)
         return true if captured_depth(stack) < @max_depth
@@ -684,7 +693,9 @@ module Bitfab
           args:,
           kwargs:,
           started_at: now,
-          span_data: without_content ? {"content_off_by_simulation_plan" => true} : nil
+          span_data: without_content ? {"content_off_by_simulation_plan" => true} : nil,
+          span_type: metadata.jev ? "llm" : nil,
+          jev_started_at: metadata.jev ? Jev.monotonic_ms : nil
         )
       end
 
@@ -708,6 +719,7 @@ module Bitfab
         return unless entry.is_a?(OpenSpan) && matching_identity?(entry.identity, event, resolved_path)
 
         stack.pop
+        record_jev_reply(entry, event.return_value) if entry.jev_started_at
         emit(entry, result: event.return_value)
         parent = stack.last
         parent.error ||= entry.error if entry.error && parent.is_a?(OpenSpan)
@@ -715,6 +727,17 @@ module Bitfab
         @drained = true if @exhausted && captured_calls_drained?
       end
 
+      def record_jev_reply(entry, reply)
+        return if entry.error || entry.jev_recorded_below
+
+        context = Jev.auto_context(reply, Jev.monotonic_ms - entry.jev_started_at)
+        return unless context
+
+        entry.contexts = [context]
+        Jev.mark_recorded(reply)
+        jev_recorded_below!
+      end
+
       def record_error(event)
         entry = @stacks[Fiber.current.object_id]&.last
         return unless entry.is_a?(OpenSpan)
@@ -849,7 +872,8 @@ module Bitfab
           lineno: event.lineno,
           name: span_name(event, method_id),
           excluded: @exclude[method_id.to_s],
-          passthrough: passthrough_wrapper?(event.parameters)
+          passthrough: passthrough_wrapper?(event.parameters),
+          jev: method_id == Jev::SYSTEM_ONE
         ).freeze
       end
 
@@ -862,6 +886,16 @@ module Bitfab
         event.callee_id || event.method_id
       end
 
+      def traced_path(event)
+        owned_path(event.path) || jev_path(event)
+      end
+
+      def jev_path(event)
+        return nil unless called_method_id(event) == Jev::SYSTEM_ONE && !sdk_frame?(event.path)
+
+        Subtree.realpath(event.path)
+      end
+
       def owned_path(path)
         @owned_paths.fetch(path) do
           resolved = Subtree.realpath(path)
@@ -945,7 +979,7 @@ module Bitfab
         return unless event.event == :call
         return if skip_boundary_call?(event)
 
-        resolved_path = owned_path(event.path)
+        resolved_path = traced_path(event)
         count_dropped_call(event, resolved_path) if resolved_path
       rescue Exception # rubocop:disable Lint/RescueException
         nil

data/lib/bitfab/traceable.rb

@@ -173,7 +173,9 @@ module Bitfab
               explicit_span_receiver: self, explicit_span_method_name: method_name,
               experiment_id:, finalize:, surface: :trace, result_recorded_elsewhere: call.recorded_elsewhere?, declared_node: true,
               on_start: call.method(:start), on_complete: call.method(:complete),
-              around_body: call.method(:run)) { super(*args, **kwargs, &block) }
+              around_body: call.method(:run)) do
+                Jev.record_node_call(method_name) { super(*args, **kwargs, &block) }
+              end
           end
         end
       end

data/lib/bitfab/version.rb

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

data/lib/bitfab.rb

@@ -30,6 +30,7 @@ require_relative "bitfab/replay_branch"
 require_relative "bitfab/client"
 require_relative "bitfab/seed"
 require_relative "bitfab/baml"
+require_relative "bitfab/jev"
 require_relative "bitfab/traceable"
 require_relative "bitfab/replay_registry"
 require_relative "bitfab/seed_cli"

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: bitfab
 version: !ruby/object:Gem::Version
-  version: 0.66.0
+  version: 0.67.0
 platform: ruby
 authors:
 - Harvest Team
@@ -175,6 +175,7 @@ files:
 - lib/bitfab/git_state.rb
 - lib/bitfab/graders.rb
 - lib/bitfab/http_client.rb
+- lib/bitfab/jev.rb
 - lib/bitfab/labels.rb
 - lib/bitfab/mock_override.rb
 - lib/bitfab/organization_members.rb