ruby_llm-contract 1.1.1 → 1.2.0

data/lib/ruby_llm/contract/concerns/usage_aggregator.rb

@@ -4,6 +4,10 @@ module RubyLLM
   module Contract
     module Concerns
       module UsageAggregator
+        # Breakdown keys an adapter may add next to input/output. They are parts
+        # of input_tokens/output_tokens, not extra tokens, so sum_tokens skips them.
+        USAGE_DETAIL_KEYS = %i[cache_read_tokens cache_write_tokens thinking_tokens].freeze
+
         private
 
         def extract_usage(trace_entry)
@@ -24,18 +28,17 @@ module RubyLLM
         end
 
         def aggregate_usage(traces)
-          input_total = 0
-          output_total = 0
+          totals = Hash.new(0)
 
           traces.each do |trace_entry|
             usage = extract_usage(trace_entry)
             next unless usage.is_a?(Hash)
 
-            input_total += usage[:input_tokens] || 0
-            output_total += usage[:output_tokens] || 0
+            [:input_tokens, :output_tokens, *USAGE_DETAIL_KEYS].each { |key| totals[key] += usage[key] || 0 }
           end
 
-          { input_tokens: input_total, output_tokens: output_total }
+          { input_tokens: totals[:input_tokens], output_tokens: totals[:output_tokens] }
+            .merge(totals.slice(*USAGE_DETAIL_KEYS).select { |_, count| count.positive? })
         end
       end
     end

data/lib/ruby_llm/contract/configuration.rb

@@ -9,13 +9,17 @@ module RubyLLM
     #
     # Then configure contract-specific options:
     #   RubyLLM::Contract.configure { |c| c.default_model = "gpt-4.1-mini" }
+    #
+    # `workflow_instrumentation = true` wraps steps and pipelines in
+    # `RubyLLM.workflow` (see WorkflowScope).
     class Configuration
-      attr_accessor :default_adapter, :default_model, :logger
+      attr_accessor :default_adapter, :default_model, :logger, :workflow_instrumentation
 
       def initialize
         @default_adapter = nil
         @default_model = nil
         @logger = nil
+        @workflow_instrumentation = false
       end
     end
   end

data/lib/ruby_llm/contract/cost_calculator.rb

@@ -7,27 +7,31 @@ module RubyLLM
     # **What this module does (public surface):**
     #
     # 1. **Fine-tune / custom-model pricing registry** — `register_model`
-    #    fills the gap left by RubyLLM 1.14's models.json: there is no
-    #    upstream `RubyLLM::Models.register` API, so fine-tuned models
-    #    (e.g. `ft:gpt-4o-custom`) need their pricing supplied locally.
+    #    fills a gap RubyLLM still has in 2.x: there is no simple
+    #    upstream API to price a model its registry lacks, so fine-tuned
+    #    models (e.g. `ft:gpt-4o-custom`) need their pricing supplied
+    #    locally. A registered price also overrides RubyLLM's for that id.
     # 2. **Lookup with fallback chain** — `calculate(model_name:, usage:)`
     #    checks the custom registry first, falls back to
     #    `RubyLLM.models.find(model_name)`, returns `nil` on miss.
     #
     # **What this module is NOT:**
     #
-    # - Not a "cost calculator" feature — the math itself
-    #   (`tokens × price_per_million / 1_000_000`) is trivial and lives
-    #   in `private_class_method :compute_cost` for internal use only.
-    # - Not a substitute for RubyLLM's pricing data — for any model in
-    #   `RubyLLM.models`, this module simply queries it.
+    # - Not a substitute for RubyLLM's pricing — for any model in
+    #   `RubyLLM.models`, RubyLLM's own `Model#cost_for` does the math.
+    # - Not the source of a finished call's cost: the RubyLLM adapter takes
+    #   that from the response (`response.cost`), which also covers what
+    #   the provider reported and RubyLLM's own retries. This module prices
+    #   pre-flight estimates and calls from adapters that report no cost.
     #
     # The reason this module exists at all is the registry + retry usage
     # aggregation across attempts (the latter sits in `Step::RetryExecutor`,
     # which calls `calculate` per attempt and sums; not in this module).
     module CostCalculator
       # Simple struct for custom-registered model pricing
-      RegisteredModel = Struct.new(:input_price_per_million, :output_price_per_million, keyword_init: true)
+      RegisteredModel = Struct.new(:input_price_per_million, :output_price_per_million,
+                                   :cache_read_price_per_million, :cache_write_price_per_million,
+                                   keyword_init: true)
 
       @custom_models = {}
 
@@ -38,13 +42,22 @@ module RubyLLM
       #   CostCalculator.register_model("ft:gpt-4o-custom",
       #     input_per_1m: 3.0, output_per_1m: 6.0)
       #
-      def self.register_model(model_name, input_per_1m:, output_per_1m:)
+      # Prompt-cache prices are optional: without `cache_read_per_1m` cache
+      # reads are charged at the input price; without `cache_write_per_1m` a
+      # call that wrote to the cache has no price (writes can cost more than
+      # input).
+      def self.register_model(model_name, input_per_1m:, output_per_1m:, cache_read_per_1m: nil,
+                              cache_write_per_1m: nil)
         validate_price!(:input_per_1m, input_per_1m)
         validate_price!(:output_per_1m, output_per_1m)
+        validate_price!(:cache_read_per_1m, cache_read_per_1m) unless cache_read_per_1m.nil?
+        validate_price!(:cache_write_per_1m, cache_write_per_1m) unless cache_write_per_1m.nil?
 
         @custom_models[model_name] = RegisteredModel.new(
           input_price_per_million: input_per_1m,
-          output_price_per_million: output_per_1m
+          output_price_per_million: output_per_1m,
+          cache_read_price_per_million: cache_read_per_1m,
+          cache_write_price_per_million: cache_write_per_1m
         )
       end
 
@@ -60,8 +73,9 @@ module RubyLLM
 
       # Look up cost for a single model + usage hash.
       # Returns nil if model is unknown (custom registry miss + RubyLLM miss),
-      # so callers can decide whether to refuse the call or proceed (see
-      # `on_unknown_pricing:` step option for the budget-gating policy).
+      # or if any used token category has no price, so callers can decide
+      # whether to refuse the call or proceed (see `on_unknown_pricing:` step
+      # option for the budget-gating policy).
       #
       #   CostCalculator.calculate(
       #     model_name: "gpt-4o-mini",
@@ -69,13 +83,17 @@ module RubyLLM
       #   )
       #   # => 0.00069 (or nil if model not registered)
       #
-      # Math is intentionally simple and private — this method is the
-      # primary public entry point. Aggregating across retry attempts is
-      # done in `Step::RetryExecutor`, not here.
-      def self.calculate(model_name:, usage:)
+      # `usage[:input_tokens]` is the whole prompt; `:cache_read_tokens` and
+      # `:cache_write_tokens`, when present, are the parts of it served from or
+      # written to the provider's prompt cache. `provider:` picks the provider's
+      # price for a model id several providers list at different prices.
+      #
+      # Aggregating across retry attempts is done in `Step::RetryExecutor`,
+      # not here.
+      def self.calculate(model_name:, usage:, provider: nil)
         return nil unless model_name && usage.is_a?(Hash)
 
-        model_info = find_model(model_name)
+        model_info = find_model(model_name, provider: provider)
         return nil unless model_info
 
         compute_cost(model_info, usage)
@@ -83,55 +101,90 @@ module RubyLLM
         nil
       end
 
-      def self.compute_cost(model_info, usage)
-        prices = prices_for(model_info)
-        return nil unless prices
-
-        input_cost = token_cost(usage[:input_tokens], prices[:input])
-        output_cost = token_cost(usage[:output_tokens], prices[:output])
-        (input_cost + output_cost).round(6)
+      def self.registered?(model_name)
+        @custom_models.key?(model_name)
       end
 
       # Two shapes reach here. `RegisteredModel` (our own struct, from
-      # `register_model`) exposes flat `*_price_per_million` readers. RubyLLM 2.0
-      # moved provider pricing into a nested value object and dropped those flat
-      # readers, so it has to be walked: pricing -> text_tokens -> standard.
+      # `register_model`) exposes flat `*_price_per_million` readers and is
+      # priced here. A RubyLLM model is priced by RubyLLM itself
+      # (`Model#cost_for`), which knows cache, long-context and per-provider
+      # prices and returns nil when a used category has no price.
       #
-      # nil means "unknown pricing". The guards keep the two known shapes
-      # explicit, but `calculate` rescues StandardError, so a future shape this
-      # walk cannot read also degrades to nil: `max_cost` then refuses the call
-      # citing "no pricing data" rather than the real cause, and paths that floor
-      # unknown cost to 0.0 (eval cost estimates, report totals) undercount.
-      def self.prices_for(model_info)
+      # nil means "unknown pricing". `calculate` rescues StandardError, so a
+      # future shape neither branch can read also degrades to nil: `max_cost`
+      # then refuses the call citing "no pricing data" rather than the real
+      # cause, and paths that floor unknown cost to 0.0 (eval cost estimates,
+      # report totals) undercount.
+      def self.compute_cost(model_info, usage)
         if model_info.respond_to?(:input_price_per_million)
-          return { input: model_info.input_price_per_million,
-                   output: model_info.output_price_per_million }
+          registered_cost(model_info, usage)
+        elsif model_info.respond_to?(:cost_for)
+          model_info.cost_for(tokens_for(usage)).total&.round(6)
         end
+      end
 
-        tier = model_info.pricing&.text_tokens&.standard if model_info.respond_to?(:pricing)
-        return unless tier.respond_to?(:input_per_million)
+      # input_tokens includes the cache reads and writes, so each part is
+      # charged once, at its own price. Without a cache-read price, reads are
+      # charged at the input price, which assumes they cost no more than
+      # uncached input (true of every provider in RubyLLM's registry, not
+      # guaranteed for a custom one). Cache writes can cost more than input
+      # (Anthropic charges 1.25x-2x), so without a cache-write price a call
+      # that wrote to the cache has no price rather than an undercount.
+      def self.registered_cost(model_info, usage)
+        cache_read = usage[:cache_read_tokens] || 0
+        cache_write = usage[:cache_write_tokens] || 0
+        read_price, write_price = registered_cache_prices(model_info)
+        return nil if cache_write.positive? && write_price.nil?
+
+        uncached = [(usage[:input_tokens] || 0) - cache_read - cache_write, 0].max
+        (token_cost(uncached, model_info.input_price_per_million) +
+          token_cost(cache_read, read_price) +
+          token_cost(cache_write, write_price.to_f) +
+          token_cost(usage[:output_tokens], model_info.output_price_per_million)).round(6)
+      end
+
+      # [read price, write price]; a model without a read price reads at its
+      # input price. Duck-typed registry entries may lack both readers.
+      def self.registered_cache_prices(model_info)
+        read = model_info.cache_read_price_per_million if model_info.respond_to?(:cache_read_price_per_million)
+        write = model_info.cache_write_price_per_million if model_info.respond_to?(:cache_write_price_per_million)
+        [read || model_info.input_price_per_million, write]
+      end
 
-        { input: tier.input_per_million, output: tier.output_per_million }
+      # RubyLLM counts uncached input apart from cache reads and writes; our
+      # usage counts them all in input_tokens.
+      def self.tokens_for(usage)
+        cache_read = usage[:cache_read_tokens] || 0
+        cache_write = usage[:cache_write_tokens] || 0
+        ::RubyLLM::Tokens.new(
+          input: [(usage[:input_tokens] || 0) - cache_read - cache_write, 0].max,
+          output: usage[:output_tokens] || 0,
+          cache_read: cache_read.positive? ? cache_read : nil,
+          cache_write: cache_write.positive? ? cache_write : nil,
+          # Priced by RubyLLM only for models with a separate reasoning price.
+          thinking: usage[:thinking_tokens]
+        )
       end
-      private_class_method :prices_for
 
       # Provider pricing is denominated per 1M tokens; divide here to get
       # the dollar cost for the actual usage count. Named constant for
       # consistency with how RubyLLM and provider docs express prices.
       TOKENS_PER_MILLION = 1_000_000.0
 
+      # `register_model` validates both prices, so neither is nil here.
       def self.token_cost(tokens, price_per_million)
-        (tokens || 0) * (price_per_million || 0) / TOKENS_PER_MILLION
+        (tokens || 0) * price_per_million / TOKENS_PER_MILLION
       end
 
-      def self.find_model(model_name)
+      def self.find_model(model_name, provider: nil)
         # Check custom registry first
         custom = @custom_models[model_name]
         return custom if custom
 
         return nil unless defined?(RubyLLM)
 
-        RubyLLM.models.find(model_name)
+        provider ? RubyLLM.models.find(model_name, provider: provider) : RubyLLM.models.find(model_name)
       rescue StandardError
         nil
       end
@@ -146,7 +199,8 @@ module RubyLLM
       # to inspect model pricing before invoking `calculate` (e.g., to short-
       # circuit estimate when the model is unknown). Exposing it removes
       # `CostCalculator.send(:find_model)` workarounds at call sites.
-      private_class_method :compute_cost, :token_cost, :validate_price!
+      private_class_method :compute_cost, :registered_cost, :registered_cache_prices, :tokens_for, :token_cost,
+                           :validate_price!
     end
   end
 end

data/lib/ruby_llm/contract/eval/eval_history.rb

@@ -69,7 +69,8 @@ module RubyLLM
 
           lines = ["#{runs.length} runs"]
           runs.last(5).each do |r|
-            lines << "  #{r[:date]} score=#{r[:score].round(2)} cost=$#{format("%.6f", r[:total_cost] || r[:cost] || 0)}"
+            cost = format("%.6f", r[:total_cost] || r[:cost] || 0)
+            lines << "  #{r[:date]} score=#{r[:score].round(2)} cost=$#{cost}#{" + unknown" if r[:cost_unknown]}"
           end
           lines.join("\n")
         end

data/lib/ruby_llm/contract/eval/model_comparison.rb

@@ -36,8 +36,12 @@ module RubyLLM
           @reports[resolve_key(candidate)]&.total_cost
         end
 
+        # A report with unpriced cases has only a subtotal, so it cannot be
+        # called the cheapest.
         def best_for(min_score: 0.0)
-          eligible = @reports.select { |_, report| report.score > 0.0 && report.score >= min_score }
+          eligible = @reports.select do |_, report|
+            report.score > 0.0 && report.score >= min_score && !cost_unknown?(report)
+          end
           return nil if eligible.empty?
 
           eligible.min_by { |_, report| report.total_cost }&.first
@@ -45,6 +49,8 @@ module RubyLLM
 
         def cost_per_point
           @reports.transform_values do |report|
+            next Float::INFINITY if cost_unknown?(report)
+
             report.score.positive? ? report.total_cost / report.score : Float::INFINITY
           end
         end
@@ -163,6 +169,10 @@ module RubyLLM
           )
         end
 
+        def cost_unknown?(report)
+          report.respond_to?(:unknown_cost_results) && report.unknown_cost_results.any?
+        end
+
         def resolve_key(candidate)
           case candidate
           when String then candidate

data/lib/ruby_llm/contract/eval/recommender.rb

@@ -51,7 +51,8 @@ module RubyLLM
               cost_per_call: cost_per_call,
               latency: report.avg_latency_ms || Float::INFINITY,
               pass_rate_ratio: report.pass_rate_ratio,
-              total_cost: report.total_cost
+              total_cost: report.total_cost,
+              cost_unknown: unknown_cost?(report)
             }
           end
         end
@@ -114,7 +115,7 @@ module RubyLLM
 
           current_label = ModelComparison.candidate_label(@current_config)
           current_report = @comparison.reports[current_label]
-          return {} unless current_report
+          return {} if current_report.nil? || unknown_cost?(current_report)
 
           current_evaluated = current_report.results.count { |r| r.step_status != CaseResult::SKIPPED_STATUS }
           current_cases = [current_evaluated, 1].max
@@ -125,8 +126,14 @@ module RubyLLM
           { per_call: diff.round(6), monthly_at: { 10_000 => (diff * 10_000).round(2) } }
         end
 
+        # A subtotal over unpriced cases is not a cost to rank by. Zero stays
+        # "unknown pricing" too: offline runs report $0 without any call.
         def cost_known?(scored_candidate)
-          scored_candidate[:cost_per_call]&.positive?
+          !scored_candidate[:cost_unknown] && scored_candidate[:cost_per_call]&.positive?
+        end
+
+        def unknown_cost?(report)
+          report.respond_to?(:unknown_cost_results) && report.unknown_cost_results.any?
         end
       end
     end

data/lib/ruby_llm/contract/eval/report_stats.rb

@@ -85,7 +85,7 @@ module RubyLLM
         def single_shot_cost
           return nil unless production_mode?
 
-          evaluated_results.sum { |r| first_attempt_cost(r) || r.cost || 0.0 }
+          evaluated_results.sum { |r| first_attempt_cost(r) || 0.0 }
         end
 
         def effective_cost
@@ -116,9 +116,14 @@ module RubyLLM
 
         private
 
+        # With attempts, only the first one's cost: falling back to the case
+        # cost when it is unpriced would report a later attempt's spend as the
+        # single-shot cost.
         def first_attempt_cost(result)
-          first = (result.attempts || []).first
-          first && first[:cost]
+          attempts = result.attempts || []
+          return result.cost if attempts.empty?
+
+          attempts.first[:cost]
         end
 
         def first_attempt_latency(result)

data/lib/ruby_llm/contract/eval/report_storage.rb

@@ -52,8 +52,10 @@ module RubyLLM
 
         private
 
+        # cost_unknown is written only when true, so older files and runs that
+        # priced every case keep the same shape.
         def history_entry
-          {
+          entry = {
             date: Time.now.strftime("%Y-%m-%d"),
             score: @stats.score,
             total_cost: @stats.total_cost,
@@ -61,6 +63,8 @@ module RubyLLM
             pass_rate_ratio: @stats.pass_rate_ratio,
             cases_count: @stats.evaluated_results_count
           }
+          entry[:cost_unknown] = true if @stats.unknown_cost_results.any?
+          entry
         end
 
         def serialize_for_baseline
@@ -74,13 +78,15 @@ module RubyLLM
         end
 
         def serialize_case(result)
-          {
+          serialized = {
             name: result.name,
             passed: result.passed?,
             score: result.score,
             details: result.details,
             cost: result.cost
           }
+          serialized[:cost_unknown] = true if result.cost_unknown?
+          serialized
         end
 
         def storage_path(root_dir, extension, model:, reasoning_effort: nil)

data/lib/ruby_llm/contract/pipeline/base.rb

@@ -48,7 +48,9 @@ module RubyLLM
           end
 
           def run(input, context: {}, timeout_ms: nil)
-            Runner.new(steps: steps, context: context, timeout_ms: timeout_ms, token_budget: token_budget).call(input)
+            WorkflowScope.workflow(name || "anonymous pipeline") do
+              Runner.new(steps: steps, context: context, timeout_ms: timeout_ms, token_budget: token_budget).call(input)
+            end
           end
 
           def test(input, responses: {}, timeout_ms: nil)

data/lib/ruby_llm/contract/pipeline/runner.rb

@@ -37,7 +37,9 @@ module RubyLLM
 
         def execute_step(step_def, execution)
           step_context = build_step_context(step_def)
-          result = step_def[:step_class].run(execution.current_input, context: step_context)
+          result = WorkflowScope.step(step_def[:alias]) do
+            step_def[:step_class].run(execution.current_input, context: step_context)
+          end
 
           execution.record_step(step_def[:alias], result)
         end
@@ -83,6 +85,7 @@ module RubyLLM
             total_usage: aggregate_usage(traces),
             step_traces: traces
           )
+          warn_incomplete_budget_usage(trace)
 
           Result.new(
             status: execution.status, step_results: execution.step_results,
@@ -91,6 +94,13 @@ module RubyLLM
           )
         end
 
+        def warn_incomplete_budget_usage(trace)
+          return unless @token_budget && !trace.usage_complete?
+
+          warn "[ruby_llm-contract] token_budget checked against incomplete token counts: " \
+               "a provider left out usage for at least one step, so the total undercounts"
+        end
+
         def elapsed_ms(start_time)
           ((Process.clock_gettime(Process::CLOCK_MONOTONIC) - start_time) * 1000).round
         end

data/lib/ruby_llm/contract/pipeline/trace.rb

@@ -40,6 +40,16 @@ module RubyLLM
           @step_traces.any? { |step_trace| step_trace.respond_to?(:cost_unknown?) && step_trace.cost_unknown? }
         end
 
+        # False when any step's provider left out token counts, so total_usage
+        # (and a token_budget checked against it) undercounts.
+        def usage_complete?
+          return true unless @step_traces.is_a?(Array)
+
+          @step_traces.none? do |step_trace|
+            step_trace.respond_to?(:usage_complete) && step_trace.usage_complete == false
+          end
+        end
+
         def to_h
           { trace_id: @trace_id, total_latency_ms: @total_latency_ms,
             total_usage: @total_usage, step_traces: @step_traces,

data/lib/ruby_llm/contract/provider_options.rb

@@ -0,0 +1,45 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Contract
+    # Options in the provider's own request vocabulary, passed to
+    # `chat.with_provider_options` and merged into the request as-is. Keys the
+    # contract itself sets are refused: an override would make the call differ
+    # from what max_cost, max_output and the schema were checked against.
+    # Only top-level keys are checked; a provider that nests these settings
+    # (Gemini's generationConfig, for one) is not.
+    module ProviderOptions
+      RESERVED_KEYS = %w[model messages input instructions system temperature stream
+                         max_tokens max_output_tokens max_completion_tokens].freeze
+      # Where OpenAI's protocols put the structured-output schema.
+      SCHEMA_KEYS = %w[response_format text].freeze
+
+      def self.validate!(options, schema: false)
+        raise ArgumentError, "provider_options must be a Hash, got #{options.class}" unless options.is_a?(Hash)
+
+        reserved = schema ? RESERVED_KEYS + SCHEMA_KEYS : RESERVED_KEYS
+        clashes = options.keys.map(&:to_s) & reserved
+        return options if clashes.empty?
+
+        raise ArgumentError, "provider_options cannot set #{clashes.join(", ")}: the step sets " \
+                             "#{clashes.one? ? "it" : "them"} (use the step DSL or context instead)"
+      end
+
+      # A frozen deep copy, so a hash the caller keeps mutating cannot change
+      # a step definition shared across threads.
+      def self.frozen_copy(object)
+        case object
+        when Hash then object.to_h { |key, value| [key, frozen_copy(value)] }.freeze
+        when Array then object.map { |value| frozen_copy(value) }.freeze
+        else object.frozen? ? object : object.dup.freeze
+        end
+      end
+
+      # Context keys win over the class-level ones; nil when nothing is set.
+      def self.merge(class_level, context_level)
+        merged = (class_level || {}).merge(context_level || {})
+        merged.empty? ? nil : merged
+      end
+    end
+  end
+end

data/lib/ruby_llm/contract/step/adapter_caller.rb

@@ -30,7 +30,8 @@ module RubyLLM
             status: :adapter_error,
             raw_output: nil,
             parsed_output: nil,
-            validation_errors: [e.message]
+            validation_errors: [e.message],
+            trace: Trace.new(error_class: e.class.name, error_type: e.class)
           )
           [result, 0]
         end