ruby_llm-contract 1.1.1 → 1.2.0

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

@@ -26,9 +26,9 @@ module RubyLLM
                              context: context).results.first
           end
 
-          def estimate_cost(input:, model: nil, attachment: nil)
+          def estimate_cost(input:, model: nil, attachment: nil, provider: nil)
             model_name = estimated_model_name(model)
-            model_info = CostCalculator.find_model(model_name)
+            model_info = CostCalculator.find_model(model_name, provider: provider)
             return nil unless model_info
 
             text_tokens = TokenEstimator.estimate(build_messages(input))
@@ -48,12 +48,13 @@ module RubyLLM
               output_tokens_estimate: output_tokens,
               estimated_cost: CostCalculator.calculate(
                 model_name: model_name,
-                usage: { input_tokens: input_tokens, output_tokens: output_tokens }
+                usage: { input_tokens: input_tokens, output_tokens: output_tokens },
+                provider: provider
               )
             }
           end
 
-          def estimate_eval_cost(eval_name, models: nil)
+          def estimate_eval_cost(eval_name, models: nil, provider: nil)
             defn = send(:all_eval_definitions)[eval_name.to_s]
             raise ArgumentError, "No eval '#{eval_name}' defined" unless defn
 
@@ -61,7 +62,7 @@ module RubyLLM
             cases = defn.build_dataset.cases
 
             model_list.each_with_object({}) do |model_name, result|
-              result[model_name] = estimate_eval_cost_for_model(cases, model_name)
+              result[model_name] = estimate_eval_cost_for_model(cases, model_name, provider)
             end
           end
 
@@ -92,7 +93,8 @@ module RubyLLM
 
           # Forwarded to the adapter as-is. A key known but not forwarded would be
           # silently dropped, so the known list is built from this one.
-          ADAPTER_CONTEXT_KEYS = %i[provider assume_model_exists max_tokens reasoning_effort attachment].freeze
+          ADAPTER_CONTEXT_KEYS = %i[provider assume_model_exists max_tokens reasoning_effort attachment
+                                    provider_options].freeze
           KNOWN_CONTEXT_KEYS = (%i[adapter model temperature retry_policy_override] + ADAPTER_CONTEXT_KEYS).freeze
 
           include Concerns::ContextHelpers
@@ -101,7 +103,7 @@ module RubyLLM
             context = safe_context(context)
             warn_unknown_context_keys(context)
 
-            result = dispatch_run(input, context)
+            result = WorkflowScope.workflow(name || "anonymous step") { dispatch_run(input, context) }
             log_result(result)
             invoke_around_call(input, result)
           end
@@ -145,9 +147,9 @@ module RubyLLM
             [0, true]
           end
 
-          def estimate_eval_cost_for_model(cases, model_name)
+          def estimate_eval_cost_for_model(cases, model_name, provider)
             cases.sum do |test_case|
-              estimate = estimate_cost(input: test_case.input, model: model_name)
+              estimate = estimate_cost(input: test_case.input, model: model_name, provider: provider)
               # Two misses floor to 0.0, not one: model absent from the registry
               # (nil estimate), or present with unreadable pricing (hash whose
               # estimated_cost is nil). Documented as a floor, not a fail-closed.
@@ -195,7 +197,7 @@ module RubyLLM
 
           def runtime_settings(context)
             policy = context.key?(:retry_policy_override) ? context[:retry_policy_override] : retry_policy
-            extra = context.slice(*ADAPTER_CONTEXT_KEYS)
+            extra = context.slice(*ADAPTER_CONTEXT_KEYS).merge(merged_provider_options(context))
 
             # Always pass the class-level `thinking` config to the adapter when
             # set, so fields like `budget` survive a per-call `reasoning_effort`
@@ -220,6 +222,13 @@ module RubyLLM
             }
           end
 
+          # Class-level provider_options with the call's merged over them; a
+          # one-key hash so `merge` drops it entirely when neither is set.
+          def merged_provider_options(context)
+            merged = ProviderOptions.merge(provider_options, context[:provider_options])
+            merged ? { provider_options: merged } : {}
+          end
+
           def current_model_config
             policy = retry_policy
             if policy && policy.config_list.any?
@@ -274,7 +283,9 @@ module RubyLLM
               on_unknown_attachment_size: on_unknown_attachment_size,
               temperature: context_temperature || temperature,
               extra_options: extra_options,
-              observers: class_observers
+              observers: class_observers,
+              on_incomplete_output: on_incomplete_output,
+              token_count: token_count
             )
           end
 
@@ -287,12 +298,18 @@ module RubyLLM
                   "model=#{trace.model} status=#{result.status} " \
                   "latency=#{trace.latency_ms}ms " \
                   "tokens=#{trace.usage&.dig(:input_tokens) || 0}+#{trace.usage&.dig(:output_tokens) || 0} " \
-                  "cost=$#{format("%.6f", trace.cost || 0)}"
+                  "cost=#{log_cost(trace)}"
             logger.info(msg)
 
             log_failed_observations(result, logger)
           end
 
+          def log_cost(trace)
+            return "unknown" if trace.respond_to?(:cost_unknown?) && trace.cost_unknown?
+
+            "$#{format("%.6f", trace.cost || 0)}"
+          end
+
           def log_failed_observations(result, logger)
             failed = result.observations.select { |o| !o[:passed] }
             return if failed.empty?

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

@@ -185,6 +185,46 @@ module RubyLLM
           inherited_value(:on_unknown_attachment_size) || UnknownPolicy::DEFAULT
         end
 
+        TOKEN_COUNT_MODES = %i[estimate exact].freeze
+        TOKEN_COUNT_DEFAULT = :estimate
+
+        # How max_input and max_cost measure the input before a call.
+        # `:estimate` (default) is the local chars/4 heuristic. `:exact` asks
+        # the provider (`chat.count_tokens`): one extra request per attempt,
+        # attachments included, no attachment_token_estimate needed. A provider
+        # or adapter that cannot count refuses the call (:limit_exceeded).
+        def token_count(mode = nil)
+          if mode
+            unless TOKEN_COUNT_MODES.include?(mode)
+              raise ArgumentError, "token_count must be one of #{TOKEN_COUNT_MODES.inspect}, got #{mode.inspect}"
+            end
+
+            return @token_count = mode
+          end
+
+          inherited_value(:token_count) || TOKEN_COUNT_DEFAULT
+        end
+
+        INCOMPLETE_OUTPUT_MODES = %i[accept refuse].freeze
+        INCOMPLETE_OUTPUT_DEFAULT = :accept
+
+        # `:refuse` fails a response the provider cut off at a token limit
+        # (:output_truncated) or stopped with a content filter
+        # (:content_filtered), before validation, so it is never :ok.
+        # `:accept` (default) validates it like any other response.
+        def on_incomplete_output(mode = nil)
+          if mode
+            unless INCOMPLETE_OUTPUT_MODES.include?(mode)
+              raise ArgumentError, "on_incomplete_output must be one of #{INCOMPLETE_OUTPUT_MODES.inspect}, " \
+                                   "got #{mode.inspect}"
+            end
+
+            return @on_incomplete_output = mode
+          end
+
+          inherited_value(:on_incomplete_output) || INCOMPLETE_OUTPUT_DEFAULT
+        end
+
         def model(name = nil)
           if name == :default
             @model = UNSET
@@ -215,6 +255,20 @@ module RubyLLM
           inherited_value_with_reset(:temperature)
         end
 
+        # Request options in the provider's vocabulary, e.g.
+        # `provider_options service_tier: "flex"`. Replaces an inherited hash;
+        # `provider_options :default` stops inheriting. `context: { provider_options: }`
+        # is merged over it per call.
+        def provider_options(options = nil)
+          if options == :default
+            @provider_options = UNSET
+            return nil
+          end
+          return @provider_options = ProviderOptions.frozen_copy(ProviderOptions.validate!(options)) if options
+
+          inherited_value_with_reset(:provider_options)
+        end
+
         def thinking(effort: nil, budget: nil)
           if effort == :default
             @thinking = UNSET

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

@@ -8,6 +8,7 @@ module RubyLLM
 
         def check_limits(messages)
           return nil unless max_input || max_cost
+          return check_counted_limits(messages) if exact_token_count?
 
           text_tokens = TokenEstimator.estimate(messages)
           attachment_tokens, attachment_error = resolve_attachment_tokens
@@ -23,6 +24,31 @@ module RubyLLM
           build_limit_result(messages, estimated, errors)
         end
 
+        # `token_count :exact`: the provider counts the input, attachments
+        # included. A count that cannot be had refuses the call rather than
+        # falling back to the heuristic under the name "exact".
+        def check_counted_limits(messages)
+          counted, count_error = count_input_tokens(messages)
+          return build_limit_result(messages, nil, [count_error], method: :exact) if count_error
+
+          errors = collect_limit_errors(counted, method: :exact)
+          errors.empty? ? nil : build_limit_result(messages, counted, errors, method: :exact)
+        end
+
+        def count_input_tokens(messages)
+          unless count_adapter.respond_to?(:count_tokens)
+            return [nil, "token_count :exact needs an adapter that counts tokens; " \
+                         "#{count_adapter.class} does not"]
+          end
+
+          counted = count_adapter.count_tokens(messages: messages, **count_options)
+          return [counted, nil] if counted.is_a?(Integer) && !counted.negative?
+
+          [nil, "token_count :exact: no usable token count (got #{counted.inspect})"]
+        rescue ::RubyLLM::Error, ::Faraday::Error => e
+          [nil, "token_count :exact: the provider could not count the input tokens (#{e.message})"]
+        end
+
         # Fail-closed: when an attachment is passed via context but no
         # attachment_token_estimate is declared, the gem cannot bound vision/
         # PDF cost. Refuses with a clear error unless on_unknown_attachment_size
@@ -50,12 +76,17 @@ module RubyLLM
           [estimate, nil]
         end
 
-        def collect_limit_errors(estimated)
+        def collect_limit_errors(estimated, method: :heuristic)
           errors = []
           if max_input && estimated > max_input
-            errors << "Input token limit exceeded: estimated #{estimated} tokens (heuristic ±30%), max #{max_input}"
+            measured = if method == :exact
+                         "#{estimated} tokens (counted by the provider)"
+                       else
+                         "estimated #{estimated} tokens (heuristic ±30%)"
+                       end
+            errors << "Input token limit exceeded: #{measured}, max #{max_input}"
           end
-          append_cost_error(estimated, errors) if max_cost
+          append_cost_error(estimated, errors, method) if max_cost
           errors
         end
 
@@ -66,18 +97,23 @@ module RubyLLM
         # for models expensive on completion side.
         DEFAULT_OUTPUT_RATIO = 1
 
-        def append_cost_error(estimated, errors)
+        def append_cost_error(estimated, errors, method = :heuristic)
           estimated_output = effective_max_output || (estimated * DEFAULT_OUTPUT_RATIO)
           estimated_cost = CostCalculator.calculate(
             model_name: model_name,
-            usage: { input_tokens: estimated, output_tokens: estimated_output }
+            usage: { input_tokens: estimated, output_tokens: estimated_output },
+            provider: pricing_provider
           )
 
           if estimated_cost.nil?
             handle_unknown_pricing(errors)
           elsif estimated_cost > max_cost
-            errors << "Cost limit exceeded: estimated $#{format("%.6f", estimated_cost)} " \
-                      "(#{estimated} input + #{estimated_output} output tokens, heuristic ±30%), " \
+            tokens = if method == :exact
+                       "#{estimated} input tokens counted by the provider + #{estimated_output} output"
+                     else
+                       "#{estimated} input + #{estimated_output} output tokens, heuristic ±30%"
+                     end
+            errors << "Cost limit exceeded: estimated $#{format("%.6f", estimated_cost)} (#{tokens}), " \
                       "max $#{format("%.6f", max_cost)}"
           end
         end
@@ -93,17 +129,16 @@ module RubyLLM
           end
         end
 
-        def build_limit_result(messages, estimated, errors)
+        # No request was sent, so usage is zero; the measured input goes
+        # alongside, and is left out when it could not be measured.
+        def build_limit_result(messages, estimated, errors, method: :heuristic)
+          usage = { input_tokens: 0, output_tokens: 0, estimated_input_tokens: estimated, estimate_method: method }
           Result.new(
             status: :limit_exceeded,
             raw_output: nil,
             parsed_output: nil,
             validation_errors: errors,
-            trace: Trace.new(
-              messages: messages, model: model_name,
-              usage: { input_tokens: 0, output_tokens: 0, estimated_input_tokens: estimated,
-                       estimate_method: :heuristic }
-            )
+            trace: Trace.new(messages: messages, model: model_name, usage: usage.compact)
           )
         end
       end

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

@@ -4,12 +4,26 @@ module RubyLLM
   module Contract
     module Step
       class ResultBuilder
-        def initialize(contract_definition:, output_type:, output_schema:, model:, observers:)
+        # Why the provider stopped, worded for a failed result. Anthropic also
+        # reports an exhausted context window as :max_tokens, so the message
+        # does not claim max_output was the limit hit.
+        FINISH_REASON_ERRORS = {
+          max_tokens: "provider reported token-limit termination",
+          content_filter: "provider reported content filtering or refusal"
+        }.freeze
+
+        # Statuses for `on_incomplete_output :refuse`.
+        INCOMPLETE_STATUSES = { max_tokens: :output_truncated, content_filter: :content_filtered }.freeze
+
+        def initialize(contract_definition:, output_type:, output_schema:, model:, observers:, max_output: nil,
+                       on_incomplete_output: Dsl::INCOMPLETE_OUTPUT_DEFAULT)
           @contract_definition = contract_definition
           @output_type = output_type
           @output_schema = output_schema
           @model = model
           @observers = observers
+          @max_output = max_output
+          @on_incomplete_output = on_incomplete_output
         end
 
         def error_result(error_result:, messages:)
@@ -18,20 +32,24 @@ module RubyLLM
             raw_output: error_result.raw_output,
             parsed_output: error_result.parsed_output,
             validation_errors: error_result.validation_errors,
-            trace: Trace.new(messages: messages, model: @model)
+            trace: Trace.new(messages: messages, model: @model, **error_fields(error_result.trace))
           )
         end
 
         def success_result(response:, messages:, latency_ms:, input:)
           raw_output = response.content
+          trace = Trace.new(messages: messages, model: @model, latency_ms: latency_ms, usage: response.usage,
+                            **response_accounting(response))
+          incomplete = incomplete_result(response, trace)
+          return incomplete if incomplete
+
           validation_result = validate_output(raw_output, input)
-          trace = Trace.new(messages: messages, model: @model, latency_ms: latency_ms, usage: response.usage)
 
           Result.new(
             status: validation_result[:status],
             raw_output: raw_output,
             parsed_output: validation_result[:parsed_output],
-            validation_errors: validation_result[:errors],
+            validation_errors: validation_errors(validation_result, response),
             trace: trace,
             observations: observations_for(validation_result, input)
           )
@@ -39,6 +57,50 @@ module RubyLLM
 
         private
 
+        def error_fields(trace)
+          return {} unless trace.respond_to?(:error_class) && trace.error_class
+
+          { error_class: trace.error_class, error_type: trace.error_type }
+        end
+
+        # Decided before validation, so neither validate blocks nor observers
+        # see a response the step refuses.
+        def incomplete_result(response, trace)
+          return unless @on_incomplete_output == :refuse
+
+          status = INCOMPLETE_STATUSES[finish_reason(response)]
+          return unless status
+
+          Result.new(status: status, raw_output: response.content, parsed_output: nil,
+                     validation_errors: [finish_reason_error(finish_reason(response))], trace: trace)
+        end
+
+        def finish_reason(response)
+          response.respond_to?(:finish_reason) ? response.finish_reason : nil
+        end
+
+        def finish_reason_error(reason)
+          message = FINISH_REASON_ERRORS[reason]
+          return message unless message && reason == :max_tokens && @max_output
+
+          "#{message}; configured max_output: #{@max_output}"
+        end
+
+        # Only a failed result says why the provider stopped; a result that
+        # passed validation keeps finish_reason in its trace alone.
+        def validation_errors(validation_result, response)
+          errors = validation_result[:errors]
+          return errors if validation_result[:status] == :ok
+
+          message = finish_reason_error(finish_reason(response))
+          message ? [*errors, message] : errors
+        end
+
+        # Custom adapters may return any object with `content` and `usage`.
+        def response_accounting(response)
+          response.respond_to?(:trace_accounting) ? response.trace_accounting : {}
+        end
+
         def observations_for(validation_result, input)
           return [] unless validation_result[:status] == :ok && @observers.any?
 

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

@@ -6,6 +6,9 @@ module RubyLLM
       module RetryExecutor
         include Concerns::UsageAggregator
 
+        # Copied from each attempt's trace into its attempts entry when present.
+        ATTEMPT_TRACE_FIELDS = %i[usage latency_ms cost finish_reason error_class].freeze
+
         private
 
         def run_with_retry(input, adapter:, default_model:, policy:, context_temperature: nil, extra_options: {})
@@ -39,7 +42,9 @@ module RubyLLM
             observations: last.observations,
             trace: last.trace.merge(
               attempts: attempt_log, usage: aggregated_usage,
-              cost: total_cost, latency_ms: total_latency
+              cost: total_cost, latency_ms: total_latency,
+              usage_complete: aggregate_flag(all_attempts, :usage_complete),
+              cost_complete: aggregate_cost_complete(all_attempts)
             )
           )
         end
@@ -53,12 +58,38 @@ module RubyLLM
         end
 
         def append_trace_fields(entry, trace)
-          entry[:usage] = trace.usage if trace.respond_to?(:usage) && trace.usage
-          entry[:latency_ms] = trace.latency_ms if trace.respond_to?(:latency_ms) && trace.latency_ms
-          entry[:cost] = trace.cost if trace.respond_to?(:cost) && trace.cost
+          ATTEMPT_TRACE_FIELDS.each do |field|
+            value = trace_value(trace, field)
+            entry[field] = value if value
+          end
+          entry[:cost_unknown] = true if trace_value(trace, :cost_unknown?)
+          entry[:usage_complete] = false if trace_value(trace, :usage_complete) == false
           entry
         end
 
+        def trace_value(trace, method)
+          trace.public_send(method) if trace.respond_to?(method)
+        end
+
+        # The subtotal is complete only if every attempt was; the last attempt's
+        # flag must not speak for it. nil when no attempt reported accounting
+        # state, so a retried Test-adapter step serializes exactly as before.
+        def aggregate_cost_complete(all_attempts)
+          traces = all_attempts.map { |a| a[:result].trace }
+          return false if traces.any? { |trace| trace.respond_to?(:cost_unknown?) && trace.cost_unknown? }
+
+          aggregate_flag(all_attempts, :cost_complete)
+        end
+
+        def aggregate_flag(all_attempts, flag)
+          flags = all_attempts.map { |a| a[:result].trace }
+                              .select { |trace| trace.respond_to?(flag) }
+                              .map { |trace| trace.public_send(flag) }.compact
+          return nil if flags.empty?
+
+          flags.all?
+        end
+
         def sum_attempt_costs(all_attempts)
           costs = extract_trace_values(all_attempts, :cost)
           return nil if costs.empty?

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

@@ -4,7 +4,7 @@ module RubyLLM
   module Contract
     module Step
       class RetryPolicy
-        attr_reader :max_attempts, :retryable_statuses
+        attr_reader :max_attempts, :retryable_statuses, :retryable_errors
 
         # Breaking (0.7.0): :adapter_error removed from defaults. ruby_llm's Faraday
         # middleware already retries transport errors (RateLimitError, ServerError,
@@ -18,6 +18,7 @@ module RubyLLM
         def initialize(models: nil, attempts: nil, retry_on: nil, &block)
           @configs = []
           @retryable_statuses = DEFAULT_RETRY_ON.dup
+          @retryable_errors = []
 
           if block
             @max_attempts = 1
@@ -49,12 +50,20 @@ module RubyLLM
           @configs
         end
 
-        def retry_on(*statuses)
-          @retryable_statuses = statuses.flatten
+        # Statuses, and exception classes for adapter errors:
+        # `retry_on :validation_failed, RubyLLM::RateLimitError` retries an
+        # adapter error only when it was a RateLimitError (or a subclass), while
+        # `:adapter_error` retries any. Replaces the defaults, as before.
+        def retry_on(*conditions)
+          @retryable_statuses, @retryable_errors = split_conditions(conditions.flatten)
         end
 
         def retryable?(result)
-          retryable_statuses.include?(result.status)
+          return true if retryable_statuses.include?(result.status)
+          return false unless result.status == :adapter_error && retryable_errors.any?
+
+          error_type = result.trace.respond_to?(:error_type) ? result.trace.error_type : nil
+          !error_type.nil? && retryable_errors.any? { |klass| error_type <= klass }
         end
 
         def config_for_attempt(attempt, default_config)
@@ -78,7 +87,24 @@ module RubyLLM
           else
             @max_attempts = attempts || 1
           end
-          @retryable_statuses = Array(retry_on).dup if retry_on
+          @retryable_statuses, @retryable_errors = split_conditions(Array(retry_on)) if retry_on
+        end
+
+        # Symbols are statuses; exception classes match adapter errors. Anything
+        # else (a String status among them) would never match, so it raises
+        # instead of quietly turning retries off.
+        def split_conditions(conditions)
+          invalid = conditions.reject { |condition| condition.is_a?(Symbol) || exception_class?(condition) }
+          unless invalid.empty?
+            raise ArgumentError,
+                  "retry_on takes status symbols and exception classes, got #{invalid.map(&:inspect).join(", ")}"
+          end
+
+          conditions.partition { |condition| condition.is_a?(Symbol) }
+        end
+
+        def exception_class?(condition)
+          condition.is_a?(Class) && condition <= Exception
         end
 
         def normalize_config(entry)

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

@@ -59,7 +59,9 @@ module RubyLLM
             output_type: @config.output_type,
             output_schema: @config.output_schema,
             model: @config.model,
-            observers: @config.observers
+            observers: @config.observers,
+            max_output: @config.effective_max_output,
+            on_incomplete_output: @config.on_incomplete_output
           )
         end
 
@@ -83,10 +85,27 @@ module RubyLLM
           @config.attachment_token_estimate
         end
 
+        def exact_token_count?
+          @config.token_count == :exact
+        end
+
+        def count_adapter
+          @config.adapter
+        end
+
+        def count_options
+          @config.adapter_options
+        end
+
         def on_unknown_attachment_size
           @config.on_unknown_attachment_size
         end
 
+        # The provider the adapter will call, so the estimate uses its price.
+        def pricing_provider
+          @config.extra_options&.dig(:provider)
+        end
+
         def attachment_present?
           opts = @config.extra_options
           !opts.nil? && !opts[:attachment].nil?

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

@@ -19,7 +19,9 @@ module RubyLLM
         :on_unknown_attachment_size,
         :temperature,
         :extra_options,
-        :observers
+        :observers,
+        :on_incomplete_output,
+        :token_count
       ) do
         # Factory with sensible defaults for optional fields. Lets callers
         # (Step::Base#run_once and tests) construct a RunnerConfig without
@@ -30,7 +32,8 @@ module RubyLLM
                        output_schema: nil, max_output: nil,
                        max_input: nil, max_cost: nil, on_unknown_pricing: UnknownPolicy::DEFAULT,
                        attachment_token_estimate: nil, on_unknown_attachment_size: UnknownPolicy::DEFAULT,
-                       temperature: nil, extra_options: {}, observers: [])
+                       temperature: nil, extra_options: {}, observers: [],
+                       on_incomplete_output: Dsl::INCOMPLETE_OUTPUT_DEFAULT, token_count: Dsl::TOKEN_COUNT_DEFAULT)
           new(
             input_type: input_type, output_type: output_type,
             prompt_block: prompt_block, contract_definition: contract_definition,
@@ -41,7 +44,7 @@ module RubyLLM
             attachment_token_estimate: attachment_token_estimate,
             on_unknown_attachment_size: on_unknown_attachment_size,
             temperature: temperature, extra_options: extra_options,
-            observers: observers
+            observers: observers, on_incomplete_output: on_incomplete_output, token_count: token_count
           )
         end