ruby_llm-contract 0.2.0

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

@@ -0,0 +1,144 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Contract
+    module Step
+      # Extracted from Base to reduce class length.
+      # DSL accessor methods for step definition (input_type, output_type, prompt, etc.).
+      module Dsl # rubocop:disable Metrics/ModuleLength
+        def input_type(type = nil)
+          return @input_type = type if type
+
+          if defined?(@input_type)
+            @input_type
+          elsif superclass.respond_to?(:input_type)
+            superclass.input_type
+          else
+            String
+          end
+        end
+
+        def output_type(type = nil)
+          return @output_type = type if type
+
+          if defined?(@output_type)
+            @output_type
+          elsif defined?(@output_schema) && @output_schema
+            RubyLLM::Contract::Types::Hash
+          elsif superclass.respond_to?(:output_type)
+            superclass.output_type
+          else
+            Hash
+          end
+        end
+
+        def output_schema(&block)
+          if block
+            require "ruby_llm/schema"
+            @output_schema = ::RubyLLM::Schema.create(&block)
+          elsif defined?(@output_schema)
+            @output_schema
+          elsif superclass.respond_to?(:output_schema)
+            superclass.output_schema
+          end
+        end
+
+        def prompt(text = nil, &block)
+          if text
+            @prompt_block = proc { user text }
+          elsif block
+            @prompt_block = block
+          elsif defined?(@prompt_block) && @prompt_block
+            @prompt_block
+          elsif superclass.respond_to?(:prompt)
+            superclass.prompt
+          else
+            raise(ArgumentError, "prompt has not been set")
+          end
+        end
+
+        def contract(&block)
+          return @contract_definition = Definition.new(&block) if block
+
+          if defined?(@contract_definition) && @contract_definition
+            @contract_definition
+          elsif superclass.respond_to?(:contract)
+            superclass.contract
+          else
+            Definition.new
+          end
+        end
+
+        def validate(description, &block)
+          (@class_validates ||= []) << Invariant.new(description, block)
+        end
+
+        def class_validates
+          own = defined?(@class_validates) ? @class_validates : []
+          inherited = superclass.respond_to?(:class_validates) ? superclass.class_validates : []
+          inherited + own
+        end
+
+        def max_output(tokens = nil)
+          if tokens
+            unless tokens.is_a?(Numeric) && tokens.positive?
+              raise ArgumentError, "max_output must be positive, got #{tokens}"
+            end
+
+            return @max_output = tokens
+          end
+
+          if defined?(@max_output)
+            @max_output
+          elsif superclass.respond_to?(:max_output)
+            superclass.max_output
+          end
+        end
+
+        def max_input(tokens = nil)
+          if tokens
+            unless tokens.is_a?(Numeric) && tokens.positive?
+              raise ArgumentError, "max_input must be positive, got #{tokens}"
+            end
+
+            return @max_input = tokens
+          end
+
+          if defined?(@max_input)
+            @max_input
+          elsif superclass.respond_to?(:max_input)
+            superclass.max_input
+          end
+        end
+
+        def max_cost(amount = nil)
+          if amount
+            unless amount.is_a?(Numeric) && amount.positive?
+              raise ArgumentError, "max_cost must be positive, got #{amount}"
+            end
+
+            return @max_cost = amount
+          end
+
+          if defined?(@max_cost)
+            @max_cost
+          elsif superclass.respond_to?(:max_cost)
+            superclass.max_cost
+          end
+        end
+
+        def retry_policy(models: nil, attempts: nil, retry_on: nil, &block)
+          if block || models || attempts
+            return @retry_policy = RetryPolicy.new(models: models, attempts: attempts, retry_on: retry_on, &block)
+          end
+
+          if defined?(@retry_policy) && @retry_policy
+            @retry_policy
+          elsif superclass.respond_to?(:retry_policy)
+            superclass.retry_policy
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,64 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Contract
+    module Step
+      # Extracted from Runner to reduce class length.
+      # Handles input token limit and cost limit checks.
+      module LimitChecker
+        private
+
+        def check_limits(messages)
+          return nil unless @max_input || @max_cost
+
+          estimated = TokenEstimator.estimate(messages)
+          errors = collect_limit_errors(estimated)
+
+          return nil if errors.empty?
+
+          build_limit_result(messages, estimated, errors)
+        end
+
+        def collect_limit_errors(estimated)
+          errors = []
+          if @max_input && estimated > @max_input
+            errors << "Input token limit exceeded: estimated #{estimated} tokens, max #{@max_input}"
+          end
+          append_cost_error(estimated, errors) if @max_cost
+          errors
+        end
+
+        def append_cost_error(estimated, errors)
+          estimated_output = @max_output || 0
+          estimated_cost = CostCalculator.calculate(
+            model_name: @model,
+            usage: { input_tokens: estimated, output_tokens: estimated_output }
+          )
+
+          if estimated_cost.nil?
+            warn "[ruby_llm-contract] max_cost is configured but model '#{@model}' " \
+                 "has no pricing data — cost limit not enforced"
+          elsif estimated_cost > @max_cost
+            errors << "Cost limit exceeded: estimated $#{format("%.6f", estimated_cost)} " \
+                      "(#{estimated} input + #{estimated_output} output tokens), " \
+                      "max $#{format("%.6f", @max_cost)}"
+          end
+        end
+
+        def build_limit_result(messages, estimated, errors)
+          Result.new(
+            status: :limit_exceeded,
+            raw_output: nil,
+            parsed_output: nil,
+            validation_errors: errors,
+            trace: Trace.new(
+              messages: messages, model: @model,
+              usage: { input_tokens: 0, output_tokens: 0, estimated_input_tokens: estimated,
+                       estimate_method: :heuristic }
+            )
+          )
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,38 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Contract
+    module Step
+      class Result
+        attr_reader :status, :raw_output, :parsed_output, :validation_errors, :trace
+
+        def initialize(status:, raw_output:, parsed_output:, validation_errors: [], trace: {})
+          @status = status
+          @raw_output = raw_output
+          @parsed_output = parsed_output
+          @validation_errors = validation_errors.freeze
+          @trace = trace.freeze
+          freeze
+        end
+
+        def ok?
+          @status == :ok
+        end
+
+        def failed?
+          @status != :ok
+        end
+
+        def to_s
+          if ok?
+            "#{@status} (#{@trace})"
+          else
+            errors = @validation_errors.first(3).join(", ")
+            errors += ", ..." if @validation_errors.size > 3
+            "#{@status}: #{errors}"
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,90 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Contract
+    module Step
+      # Extracted from Base to reduce class length.
+      # Handles retry logic: run_with_retry, build_retry_result, aggregate usage, build attempt entries.
+      module RetryExecutor
+        private
+
+        def run_with_retry(input, adapter:, default_model:, policy:)
+          all_attempts = []
+
+          policy.max_attempts.times do |attempt_index|
+            model = policy.model_for_attempt(attempt_index, default_model)
+            result = run_once(input, adapter: adapter, model: model)
+            all_attempts << { attempt: attempt_index + 1, model: model, result: result }
+            break unless policy.retryable?(result)
+          end
+
+          build_retry_result(all_attempts)
+        end
+
+        def build_retry_result(all_attempts)
+          last = all_attempts.last[:result]
+          attempt_log = all_attempts.map { |attempt| build_attempt_entry(attempt) }
+          aggregated_usage = aggregate_retry_usage(all_attempts)
+          total_cost = sum_attempt_costs(all_attempts)
+          total_latency = sum_attempt_latency(all_attempts)
+
+          Result.new(
+            status: last.status, raw_output: last.raw_output,
+            parsed_output: last.parsed_output, validation_errors: last.validation_errors,
+            trace: last.trace.merge(
+              attempts: attempt_log, usage: aggregated_usage,
+              cost: total_cost, latency_ms: total_latency
+            )
+          )
+        end
+
+        def build_attempt_entry(attempt)
+          trace = attempt[:result].trace
+          entry = { attempt: attempt[:attempt], model: attempt[:model], status: attempt[:result].status }
+          append_trace_fields(entry, trace)
+        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
+          entry
+        end
+
+        def sum_attempt_costs(all_attempts)
+          costs = extract_trace_values(all_attempts, :cost)
+          return nil if costs.empty?
+
+          costs.sum.round(6)
+        end
+
+        def sum_attempt_latency(all_attempts)
+          latencies = extract_trace_values(all_attempts, :latency_ms)
+          return nil if latencies.empty?
+
+          latencies.sum
+        end
+
+        def extract_trace_values(all_attempts, method)
+          all_attempts.filter_map do |a|
+            trace = a[:result].trace
+            trace.respond_to?(method) && trace.public_send(method)
+          end
+        end
+
+        def aggregate_retry_usage(all_attempts)
+          totals = { input_tokens: 0, output_tokens: 0 }
+          all_attempts.each do |attempt|
+            usage = attempt[:result].trace
+            usage = usage.respond_to?(:usage) ? usage.usage : nil
+            next unless usage.is_a?(Hash)
+
+            totals[:input_tokens] += usage[:input_tokens] || 0
+            totals[:output_tokens] += usage[:output_tokens] || 0
+          end
+          totals
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,76 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Contract
+    module Step
+      class RetryPolicy
+        attr_reader :max_attempts, :retryable_statuses
+
+        DEFAULT_RETRY_ON = %i[validation_failed parse_error adapter_error].freeze
+
+        def initialize(models: nil, attempts: nil, retry_on: nil, &block)
+          @models = []
+          @retryable_statuses = DEFAULT_RETRY_ON.dup
+
+          if block
+            @max_attempts = 1
+            instance_eval(&block)
+          else
+            apply_keywords(models: models, attempts: attempts, retry_on: retry_on)
+          end
+
+          validate_max_attempts!
+        end
+
+        def attempts(count)
+          @max_attempts = count
+          validate_max_attempts!
+        end
+
+        def escalate(*model_list)
+          @models = model_list.flatten
+          @max_attempts = @models.length if @max_attempts < @models.length
+        end
+        alias models escalate
+
+        def model_list
+          @models
+        end
+
+        def retry_on(*statuses)
+          @retryable_statuses = statuses
+        end
+
+        def retryable?(result)
+          retryable_statuses.include?(result.status)
+        end
+
+        def model_for_attempt(attempt, default_model)
+          if @models.any?
+            @models[attempt] || @models.last
+          else
+            default_model
+          end
+        end
+
+        private
+
+        def apply_keywords(models:, attempts:, retry_on:)
+          if models
+            @models = Array(models).dup.freeze
+            @max_attempts = @models.length
+          else
+            @max_attempts = attempts || 1
+          end
+          @retryable_statuses = Array(retry_on).dup if retry_on
+        end
+
+        def validate_max_attempts!
+          return if @max_attempts.is_a?(Integer) && @max_attempts >= 1
+
+          raise ArgumentError, "attempts must be at least 1, got #{@max_attempts.inspect}"
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,126 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Contract
+    module Step
+      class Runner
+        include LimitChecker
+
+        def initialize(input_type:, output_type:, prompt_block:, contract_definition:,
+                       adapter:, model:, output_schema: nil, max_output: nil,
+                       max_input: nil, max_cost: nil)
+          @input_type = input_type
+          @output_type = output_type
+          @prompt_block = prompt_block
+          @contract_definition = contract_definition
+          @adapter = adapter
+          @model = model
+          @output_schema = output_schema
+          @max_output = max_output
+          @max_input = max_input
+          @max_cost = max_cost
+        end
+
+        def call(input)
+          validated_input = validate_input(input)
+          return validated_input if validated_input.is_a?(Result)
+
+          messages = build_and_render_prompt(input)
+        rescue RubyLLM::Contract::Error => e
+          Result.new(status: :input_error, raw_output: nil, parsed_output: nil,
+                     validation_errors: [e.message])
+        else
+          execute_pipeline(messages, input)
+        end
+
+        private
+
+        def execute_pipeline(messages, input)
+          limit_result = check_limits(messages)
+          return limit_result if limit_result
+
+          response, latency_ms = execute_adapter(messages)
+          return build_error_result(response, messages) if response.is_a?(Result)
+
+          build_result(response, messages, latency_ms, input)
+        end
+
+        def validate_input(input)
+          type = @input_type
+          if type.is_a?(Class) && !type.respond_to?(:[])
+            raise TypeError, "#{input.inspect} is not a #{type}" unless input.is_a?(type)
+          else
+            type[input]
+          end
+          nil
+        rescue Dry::Types::CoercionError, TypeError, ArgumentError => e
+          Result.new(status: :input_error, raw_output: nil, parsed_output: nil, validation_errors: [e.message])
+        end
+
+        def build_and_render_prompt(input)
+          dynamic = @prompt_block.arity >= 1
+          ast = Prompt::Builder.build(input: dynamic ? input : nil, &@prompt_block)
+
+          Prompt::Renderer.render(ast, variables: dynamic ? {} : template_variables_for(input))
+        rescue StandardError => e
+          raise RubyLLM::Contract::Error, "Prompt build failed: #{e.class}: #{e.message}"
+        end
+
+        def template_variables_for(input)
+          base = { input: input }
+          input.is_a?(Hash) ? base.merge(input.transform_keys(&:to_sym)) : base
+        end
+
+        def execute_adapter(messages)
+          start_time = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+          response = @adapter.call(messages: messages, **build_adapter_options)
+          latency_ms = ((Process.clock_gettime(Process::CLOCK_MONOTONIC) - start_time) * 1000).round
+          [response, latency_ms]
+        rescue StandardError => e
+          [Result.new(status: :adapter_error, raw_output: nil, parsed_output: nil, validation_errors: [e.message]), 0]
+        end
+
+        def build_adapter_options
+          { model: @model }.tap do |opts|
+            opts[:schema] = @output_schema if @output_schema
+            opts[:max_tokens] = @max_output if @max_output
+          end
+        end
+
+        def build_error_result(error_result, messages)
+          Result.new(
+            status: error_result.status,
+            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)
+          )
+        end
+
+        def build_result(response, messages, latency_ms, input)
+          raw_output = response.content
+          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],
+            trace: trace
+          )
+        end
+
+        def validate_output(raw_output, input)
+          Validator.validate(
+            raw_output: raw_output,
+            definition: @contract_definition,
+            output_type: @output_type,
+            input: input,
+            schema: @output_schema
+          )
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,70 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Contract
+    module Step
+      class Trace
+        include Concerns::TraceEquality
+
+        attr_reader :messages, :model, :latency_ms, :usage, :attempts, :cost
+
+        def initialize(messages: nil, model: nil, latency_ms: nil, usage: nil, attempts: nil, cost: nil)
+          @messages = messages
+          @model = model
+          @latency_ms = latency_ms
+          @usage = usage
+          @attempts = attempts
+          @cost = cost || CostCalculator.calculate(model_name: model, usage: usage)
+          freeze
+        end
+
+        KNOWN_KEYS = %i[messages model latency_ms usage attempts cost].freeze
+
+        def [](key)
+          return nil unless KNOWN_KEYS.include?(key.to_sym)
+
+          public_send(key)
+        end
+
+        def key?(key)
+          KNOWN_KEYS.include?(key.to_sym) && !public_send(key).nil?
+        end
+        alias has_key? key?
+
+        def merge(**overrides)
+          self.class.new(
+            messages: overrides.fetch(:messages, @messages),
+            model: overrides.fetch(:model, @model),
+            latency_ms: overrides.fetch(:latency_ms, @latency_ms),
+            usage: overrides.fetch(:usage, @usage),
+            attempts: overrides.fetch(:attempts, @attempts),
+            cost: overrides.fetch(:cost, @cost)
+          )
+        end
+
+        def to_h
+          { messages: @messages, model: @model, latency_ms: @latency_ms,
+            usage: @usage, attempts: @attempts, cost: @cost }.compact
+        end
+
+        def to_s
+          build_summary_parts.join(" ")
+        end
+
+        private
+
+        def build_summary_parts
+          parts = [@model || "no-model"]
+          parts << "#{@latency_ms}ms" if @latency_ms
+          parts << format_token_usage if @usage.is_a?(Hash)
+          parts << "$#{format("%.6f", @cost)}" if @cost
+          parts
+        end
+
+        def format_token_usage
+          "#{@usage[:input_tokens] || 0}+#{@usage[:output_tokens] || 0} tokens"
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/contract/step.rb

@@ -0,0 +1,10 @@
+# frozen_string_literal: true
+
+require_relative "step/trace"
+require_relative "step/result"
+require_relative "step/limit_checker"
+require_relative "step/runner"
+require_relative "step/retry_policy"
+require_relative "step/retry_executor"
+require_relative "step/dsl"
+require_relative "step/base"

data/lib/ruby_llm/contract/token_estimator.rb

@@ -0,0 +1,19 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Contract
+    module TokenEstimator
+      # Heuristic: ~4 characters per token for English text.
+      # This is a rough estimate — actual tokenization varies by model and content.
+      # Intentionally conservative (overestimates slightly) to avoid surprise costs.
+      CHARS_PER_TOKEN = 4
+
+      def self.estimate(messages)
+        return 0 unless messages.is_a?(Array)
+
+        total_chars = messages.sum { |m| m[:content].to_s.length }
+        (total_chars.to_f / CHARS_PER_TOKEN).ceil
+      end
+    end
+  end
+end

data/lib/ruby_llm/contract/types.rb

@@ -0,0 +1,11 @@
+# frozen_string_literal: true
+
+require "dry-types"
+
+module RubyLLM
+  module Contract
+    module Types
+      include Dry.Types()
+    end
+  end
+end

data/lib/ruby_llm/contract/version.rb

@@ -0,0 +1,7 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Contract
+    VERSION = "0.2.0"
+  end
+end