ruby-laya 0.1.0

data/lib/laya/router.rb

@@ -0,0 +1,353 @@
+# frozen_string_literal: true
+
+require "monitor"
+
+module Laya
+  # The upstream checkpoint each name refers to. The gem runs the ONNX export of whichever one is
+  # chosen, so a decision reports the checkpoint's own id: the weights are the same.
+  BUNDLE_REPO = Checkpoints::BUNDLE_REPO
+  DEFAULT_MODELS = {
+    "english" => [BUNDLE_REPO, nil].freeze,
+    "multilingual" => [BUNDLE_REPO, "multilingual"].freeze,
+    "typed-decisions" => [BUNDLE_REPO, "typed-decisions"].freeze
+  }.freeze
+
+  # The same checkpoints in their own repositories, for anyone who prefers them.
+  STANDALONE_MODELS = {
+    "english" => "convaiinnovations/laya",
+    "multilingual" => "convaiinnovations/laya-multilingual",
+    "typed-decisions" => "convaiinnovations/laya-typed-decisions"
+  }.freeze
+
+  MODEL_ALIASES = Checkpoints::ALIASES
+
+  # The question-id signatures of the four typed-decisions workflows, used only when
+  # `auto_task_detection` is on.
+  TYPED_DECISION_WORKFLOWS = {
+    "agent_trace_observability" => Set.new(%w[action needs_review outcome risk urgency]).freeze,
+    "customer_service" => Set.new(%w[action category churn_risk needs_human urgency]).freeze,
+    "invoice_processing" => Set.new(%w[discrepancy_severity disposition duplicate matches_order urgency]).freeze,
+    "security_incidents" => Set.new(%w[credential_compromise disposition severity true_positive urgency]).freeze
+  }.freeze
+
+  # Language subtags that mean "the English checkpoint can read this". Routing needs one bit, not
+  # a language id, so every other code resolves to the multilingual checkpoint.
+  ENGLISH_SUBTAGS = %w[en eng english].freeze
+
+  class << self
+    def normalise_name(name) = Checkpoints.normalise(name)
+    def normalize_name(name) = Checkpoints.normalise(name)
+    def repo_str(spec) = Checkpoints.repo_str(spec)
+
+    def split_model_spec(spec)
+      spec.is_a?(Array) ? [spec[0], spec[1]] : [spec, nil]
+    end
+
+    # The workflow whose question ids these are, or nil. An exact match is required, so a schema
+    # that merely contains "urgency" is never captured.
+    def match_typed_decisions_workflow(questions)
+      ids = Set.new((questions || {}).keys.map(&:to_s))
+      TYPED_DECISION_WORKFLOWS.find { |_name, signature| ids == signature }&.first
+    end
+
+    # Whether a language code says the English checkpoint can read the text: true, false, or nil
+    # when the code identifies nothing.
+    #
+    # Accepts `"en"`, `"EN"`, `"en-US"`, the POSIX `"en_US"` and `"en_US.UTF-8"`. Nil is not a
+    # verdict but the absence of one, which is what lets a language identifier abstain and the
+    # router fall through to its own detection. It is deliberately not a predicate: a `?` method
+    # that answers nil is a trap for the caller, and for anything that "simplifies" it later.
+    def english_language_hint(value)
+      return nil if value.nil?
+
+      code = value.to_s.strip.downcase.split(".", 2).first.to_s
+      primary = code.tr("_", "-").split("-", 2).first.to_s
+      return nil if primary.empty?
+
+      ENGLISH_SUBTAGS.include?(primary)
+    end
+  end
+
+  # Why a request went to the checkpoint it did.
+  class RouteDecision
+    attr_reader :model, :repo, :reason, :detection, :workflow
+
+    def initialize(model:, repo:, reason:, detection: nil, workflow: nil)
+      @model = model
+      @repo = repo
+      @reason = reason
+      @detection = detection
+      @workflow = workflow
+    end
+
+    # The payload upstream's Python puts under "routing".
+    def to_h
+      { "model" => model, "repo" => repo, "reason" => reason,
+        "detection" => detection, "workflow" => workflow }
+    end
+
+    def to_json(*) = to_h.to_json(*)
+
+    def inspect
+      "#<Laya::RouteDecision #{model.inspect} #{reason.inspect}>"
+    end
+  end
+
+  # Sends each request to the checkpoint best suited to it.
+  #
+  #   router = Laya::Router.new
+  #   router.predict({ "message" => "Mein Konto wurde zweimal belastet" }, questions) # multilingual
+  #   router.predict({ "message" => "I was charged twice" }, questions)               # english
+  #   router.predict(state, questions, model: "typed-decisions")                      # explicit
+  #
+  # Checkpoints are downloaded and built on first use, and `max_loaded` caps how many stay
+  # resident. The default of two is what automatic routing needs: it only ever chooses between
+  # english and multilingual, and a cap of one would rebuild the checkpoint it just evicted on
+  # every script switch. Raise it to three, or preload, when `typed-decisions` is also in play.
+  #
+  #   router = Laya::Router.new(preload: true)          # everything resident, routing is free
+  #   router.preload(["english", "multilingual"])       # or just the two you serve
+  #   router.attach("english", existing_agent)          # reuse an agent you already built
+  #   router.unload                                     # free memory
+  class Router
+    DEFAULT_MAX_LOADED = 2
+
+    attr_reader :models, :device, :providers, :token, :default, :auto_task_detection, :lang_guess
+    attr_accessor :max_loaded
+
+    # Builds an {Agent}. Injectable so a test, or an app with its own loading rules, can decide
+    # how a checkpoint comes into being.
+    DEFAULT_AGENT_FACTORY = lambda do |repo, subfolder: nil, **options|
+      Agent.new(repo, subfolder: subfolder, **options)
+    end
+
+    def self.open(**)
+      router = new(**)
+      return router unless block_given?
+
+      begin
+        yield router
+      ensure
+        router.close
+      end
+    end
+
+    def initialize(models: nil, device: nil, providers: nil, token: nil, max_loaded: DEFAULT_MAX_LOADED,
+                   default: "english", auto_task_detection: false, standalone_repos: false,
+                   preload: false, lang_guess: nil, threads: nil, agent_factory: nil)
+      @models = (standalone_repos ? STANDALONE_MODELS : DEFAULT_MODELS).dup
+      models&.each { |name, spec| @models[Checkpoints.normalise(name)] = spec }
+      @device = device
+      @providers = providers
+      @threads = threads
+      @token = token || Hub.token
+      @max_loaded = [1, Integer(max_loaded)].max
+      @default = Checkpoints.normalise(default)
+      @auto_task_detection = auto_task_detection ? true : false
+      # An opt-in hint applied to every request: a language code, or a callable taking the state
+      # and returning one (or nil to abstain). Checked before the built-in detection, never
+      # before an explicit model, task or lang. This is the seam for a real language model.
+      @lang_guess = lang_guess
+      @agent_factory = agent_factory || DEFAULT_AGENT_FACTORY
+      @agents = {}
+      @order = [] # least recently used first
+      # Guards the model lifecycle and the LRU bookkeeping. Inference deliberately runs outside
+      # it, so concurrent requests share a checkpoint instead of queueing.
+      @lock = Monitor.new
+      self.preload if preload
+    end
+
+    # ---------------------------------------------------------------- loading
+
+    # The agent for `name`, built on first use. Concurrent callers share one.
+    def load(name)
+      key = Checkpoints.normalise(name)
+      @lock.synchronize do
+        if @agents.key?(key)
+          touch(key)
+          next @agents[key]
+        end
+
+        repo, subfolder = Laya.split_model_spec(@models.fetch(key))
+        agent = @agent_factory.call(repo, subfolder: subfolder, device: @device, providers: @providers,
+                                          token: @token, threads: @threads)
+        @agents[key] = agent
+        @order << key
+        evict
+        agent
+      end
+    end
+
+    # Register an already-built agent instead of loading a second copy.
+    def attach(name, agent)
+      key = Checkpoints.normalise(name)
+      @lock.synchronize do
+        @agents[key] = agent
+        touch(key)
+        @max_loaded = [@max_loaded, @agents.length].max
+      end
+      agent
+    end
+
+    # Build checkpoints up front, so no request pays a cold load. `max_loaded` grows to fit both
+    # what is requested and what is already resident, so preloading never evicts.
+    def preload(names = nil)
+      names = (names || @models.keys).map { |name| Checkpoints.normalise(name) }
+      @lock.synchronize do
+        @max_loaded = [@max_loaded, (names | @agents.keys).length].max
+        names.each { |name| load(name) unless @agents.key?(name) }
+      end
+      self
+    end
+
+    # Free one checkpoint, or all of them.
+    def unload(name = nil)
+      @lock.synchronize do
+        keys = name ? [Checkpoints.normalise(name)] : @agents.keys
+        keys.each do |key|
+          agent = @agents.delete(key)
+          @order.delete(key)
+          agent&.close if agent.respond_to?(:close)
+        end
+      end
+      nil
+    end
+    alias close unload
+
+    # The resident checkpoints, least recently used first.
+    def loaded
+      @lock.synchronize { @order.dup }
+    end
+
+    def agents
+      @lock.synchronize { @agents.dup }
+    end
+
+    # ---------------------------------------------------------------- routing
+
+    # Decide which checkpoint to use, without loading or running anything.
+    #
+    # Precedence: explicit `model`, explicit `task`, a detected workflow (opt-in), explicit
+    # `lang`, a `lang_guess` hint, detected script and language, then the default.
+    def route(state, questions = nil, model: nil, task: nil, lang: nil, lang_guess: nil)
+      return decide(model, "explicit model=#{model.inspect}") unless model.nil?
+      return decide(task_name(task), "explicit task=#{task.inspect}") unless task.nil?
+
+      workflow = Laya.match_typed_decisions_workflow(questions || {})
+      if workflow && auto_task_detection
+        return decide("typed-decisions",
+                      "question ids match the #{workflow.inspect} typed-decisions workflow",
+                      workflow: workflow)
+      end
+      unless lang.nil?
+        return decide(checkpoint_for(Laya.english_language_hint(lang)), "explicit lang=#{lang.inspect}",
+                      workflow: workflow)
+      end
+
+      hinted = hinted_decision(state, lang_guess, workflow)
+      return hinted if hinted
+
+      detected(state, workflow)
+    end
+
+    # ---------------------------------------------------------------- running
+
+    # Route, then answer every question on the chosen checkpoint. The {Result} carries the
+    # decision that was made.
+    def predict(state, questions, model: nil, task: nil, lang: nil, lang_guess: nil)
+      decision = route(state, questions, model: model, task: task, lang: lang, lang_guess: lang_guess)
+      load(decision.model).predict(state, questions).with_routing(decision)
+    end
+    alias system_one predict
+
+    def inspect
+      "#<Laya::Router loaded=#{loaded.inspect} max_loaded=#{max_loaded} default=#{default.inspect}>"
+    end
+
+    private
+
+    def decide(name, reason, detection: nil, workflow: nil)
+      key = Checkpoints.normalise(name)
+      RouteDecision.new(model: key, repo: Checkpoints.repo_str(@models.fetch(key)), reason: reason,
+                        detection: detection, workflow: workflow)
+    end
+
+    def task_name(task)
+      task.to_s.downcase.tr("-", "_") == "typed_decisions" ? "typed-decisions" : task
+    end
+
+    def checkpoint_for(english)
+      english ? "english" : "multilingual"
+    end
+
+    # A caller's hint, per call first and then the one installed on the router. Only a hint that
+    # actually answers the question routes; anything else falls through to detection.
+    def hinted_decision(state, per_call, workflow)
+      [["lang_guess", per_call], ["Router(lang_guess=...)", lang_guess]].each do |source, hint|
+        english = resolve_hint(hint, state)
+        next if english.nil?
+
+        return decide(checkpoint_for(english),
+                      "#{source}: the caller identified this as #{english ? 'English' : 'non-English'} text",
+                      workflow: workflow)
+      end
+      nil
+    end
+
+    def resolve_hint(hint, state)
+      return nil if hint.nil?
+
+      Laya.english_language_hint(hint.respond_to?(:call) ? hint.call(state) : hint)
+    end
+
+    def detected(state, workflow)
+      detection = Lang.analyse(state)
+      name, reason = read_detection(detection)
+      decide(name, reason, detection: detection, workflow: workflow)
+    end
+
+    def read_detection(detection)
+      if detection["script"] == "unknown"
+        [default, "no letters detected in state; using default (#{default})"]
+      elsif detection["script"] != "latin"
+        ["multilingual", format("non-Latin script (%s, %.0f%% of letters); the English checkpoint " \
+                                "cannot read it", detection["script"], 100 * detection["non_latin_fraction"])]
+      elsif !detection["is_english"]
+        ["multilingual", non_english_reason(detection)]
+      elsif detection["language_undecided"]
+        # Nothing identifies the language: too short, or only content words. That is no evidence
+        # of English either, so it takes the same default as a state with no letters at all.
+        [default, "Latin script, language not identified and no non-English letters; " \
+                  "using default (#{default})"]
+      else
+        ["english", "English Latin text"]
+      end
+    end
+
+    def non_english_reason(detection)
+      if detection["language"]
+        return "Latin script but language looks like #{detection['language'].inspect}, not English"
+      end
+
+      # An unidentified Latin-script language, routed on its non-English letters alone because no
+      # stopword list here covers it.
+      format("Latin script, language not identified but %.0f%% non-English letters; not safe for " \
+             "the English checkpoint", 100 * detection["diacritic_rate"])
+    end
+
+    def touch(key)
+      @order.delete(key)
+      @order << key
+    end
+
+    def evict
+      while @order.length > @max_loaded
+        agent = @agents.delete(@order.shift)
+        agent&.close if agent.respond_to?(:close)
+      end
+      return unless @order.length < @agents.length
+
+      (@agents.keys - @order).each { |key| @agents.delete(key)&.then { |a| a.close if a.respond_to?(:close) } }
+    end
+  end
+end

data/lib/laya/runtime.rb

@@ -0,0 +1,79 @@
+# frozen_string_literal: true
+
+require "onnxruntime"
+
+module Laya
+  # The ONNX Runtime session behind an {Agent}: the published decision model, loaded once.
+  #
+  # The graph exposes three outputs and ONNX Runtime computes only the ones asked for, so
+  # answering questions never pays for the encoder states the embedding shortlist needs.
+  class Runtime
+    LOGITS = "logits"
+    ACT_LOGITS = "act_logits"
+    HIDDEN = "last_hidden_state"
+    INPUTS = %w[input_ids attention_mask marker_pos marker_mask qtype].freeze
+
+    # Execution providers per device name, most specific first. An unavailable provider is
+    # skipped by ONNX Runtime with a warning, so naming one is never fatal.
+    DEVICES = {
+      "cpu" => ["CPUExecutionProvider"],
+      "coreml" => %w[CoreMLExecutionProvider CPUExecutionProvider],
+      "mps" => %w[CoreMLExecutionProvider CPUExecutionProvider],
+      "cuda" => %w[CUDAExecutionProvider CPUExecutionProvider],
+      "gpu" => %w[CUDAExecutionProvider CPUExecutionProvider],
+      "tensorrt" => %w[TensorrtExecutionProvider CUDAExecutionProvider CPUExecutionProvider],
+      "directml" => %w[DmlExecutionProvider CPUExecutionProvider]
+    }.freeze
+
+    # Resolve `device:` / `providers:` into a provider list.
+    def self.providers_for(device: nil, providers: nil)
+      return Array(providers) unless providers.nil?
+      return DEVICES.fetch("cpu") if device.nil?
+
+      DEVICES.fetch(device.to_s.downcase.split(":").first) do
+        raise ArgumentError, "unknown device #{device.inspect}; use one of #{DEVICES.keys} " \
+                             "or pass providers: [\"...ExecutionProvider\"]"
+      end
+    end
+
+    attr_reader :path, :providers
+
+    def initialize(path, providers: nil, device: nil, threads: nil)
+      @path = path
+      @providers = Runtime.providers_for(device: device, providers: providers)
+      @session = OnnxRuntime::InferenceSession.new(
+        path, providers: @providers, intra_op_num_threads: threads
+      )
+      @mutex = Mutex.new
+    end
+
+    # Score a batch of questions. Returns `[logits, act_logits]` as nested Arrays of Float.
+    def decide(batch)
+      run([LOGITS, ACT_LOGITS], batch)
+    end
+
+    # Encoder states for a tokenized batch, `[batch, seq, hidden]`.
+    def encode(batch)
+      run([HIDDEN], batch).first
+    end
+
+    def close
+      @session = nil
+    end
+
+    def closed?
+      @session.nil?
+    end
+
+    private
+
+    def run(outputs, batch)
+      raise Error, "this agent has been closed" if closed?
+
+      feed = INPUTS.to_h { |name| [name, batch.fetch(name.to_sym)] }
+      # ONNX Runtime sessions are thread-safe for inference, but the gem's FFI pointers are not
+      # reentrant, so one request at a time per session.
+      @mutex.synchronize { @session.run(outputs, feed) }
+    end
+  end
+end

data/lib/laya/shortlist.rb

@@ -0,0 +1,197 @@
+# frozen_string_literal: true
+
+module Laya
+  # An opt-in embedding shortlist for choice questions with many labels.
+  #
+  # Options share one `head_max_len` budget, so a large label set leaves only a few tokens per
+  # label and they stop being distinguishable. {predict_shortlist} embeds the state and each
+  # option with a caller-supplied `embed_fn`, keeps the top `k`, and runs one prediction on that
+  # reduced set. `predict` itself is untouched: it still scores every criterion it is given.
+  module Shortlist
+    DEFAULT_SHORTLIST_K = 20
+
+    module_function
+
+    # The top `k` labels for `state`.
+    #
+    # `embed_fn` maps an Array of Strings to one vector per string. It is called once, with the
+    # query first and then the options in criteria order, rendered as the model would see them.
+    # When `k` is at least the number of labels every label is returned in its original order and
+    # `embed_fn` is never called. Ties keep the earlier label, and a zero vector never outranks
+    # a label that came before it.
+    def shortlist_choice(state, criteria, embed_fn, k: DEFAULT_SHORTLIST_K, instructions: nil)
+      rank(state, criteria, embed_fn, k, instructions).labels
+    end
+
+    # Shortlist every choice question, then predict once.
+    #
+    # Questions that are not choices, and choices with `k` labels or fewer, are passed through
+    # untouched. The caller's Hash is never mutated. The result carries a `shortlist` entry
+    # recording, per question, the labels kept in rank order, their cosine scores, `k`, the
+    # original label count `n`, and whether it was a pass-through.
+    def predict_shortlist(agent, state, questions, embed_fn, k: DEFAULT_SHORTLIST_K, **predict_options)
+      raise TypeError, "questions must be a Hash of question id => definition" unless questions.is_a?(Hash)
+
+      checked = Util.positive_int!(k, "k")
+      reduced = {}
+      meta = {}
+      questions.each do |id, definition|
+        unless choice?(definition)
+          reduced[id] = definition
+          next
+        end
+
+        ranking = rank_question(id, definition, state, embed_fn, checked)
+        meta[id] = ranking.to_h
+        reduced[id] = ranking.passthrough ? definition : narrow(definition, ranking.labels)
+      end
+
+      attach(predict_with(agent, state, reduced, **predict_options), meta)
+    end
+
+    # An `embed_fn` backed by the checkpoint already loaded on `agent`.
+    #
+    # It mean-pools the encoder, which costs nothing extra but is not a retriever. Measured on
+    # Banking77's 77 labels it keeps the right label in the top 20 about as often as picking 20
+    # labels at random would: mean-pooled states of this encoder sit within a couple of hundredths
+    # of each other in cosine, so the ranking carries little signal. Centering the batch widens the
+    # spread without improving recall.
+    #
+    # Pass a real bi-encoder as `embed_fn` when the shortlist has to be right. This helper is a
+    # starting point for callers who have nothing else loaded, and worth measuring on your own
+    # labels before relying on it.
+    def embed_fn_from_agent(agent, max_length: nil, batch_size: 32)
+      Util.positive_int!(max_length, "max_length") unless max_length.nil?
+      Util.positive_int!(batch_size, "batch_size")
+
+      ->(texts) { agent.embed(texts, max_length: max_length, batch_size: batch_size) }
+    end
+
+    # What ranking a question produced.
+    Ranking = Struct.new(:labels, :scores, :passthrough, :n, :k, keyword_init: true) do
+      def to_h
+        { "labels" => labels.dup, "scores" => scores, "k" => k, "n" => n, "passthrough" => passthrough }
+      end
+    end
+
+    def choice?(definition)
+      definition.is_a?(Hash) && Util.get(definition, "type").to_s == "choice"
+    end
+
+    def rank_question(id, definition, state, embed_fn, k)
+      unless Util.key?(definition, "criteria")
+        raise ArgumentError, "question #{id.inspect} is a choice but has no criteria"
+      end
+
+      rank(state, Util.get(definition, "criteria"), embed_fn, k, Util.get(definition, "instructions"))
+    end
+
+    def rank(state, criteria, embed_fn, k, instructions)
+      checked = Util.positive_int!(k, "k")
+      items = criteria_items(criteria)
+      labels = items.map(&:first)
+      if checked >= items.length
+        return Ranking.new(labels: labels, scores: nil, passthrough: true, n: items.length, k: checked)
+      end
+
+      matrix = embeddings(embed_fn, [query_text(state, instructions)] + option_texts(items))
+      scores = cosine(matrix.first, matrix.drop(1))
+      order = (0...items.length).sort_by { |i| [-scores[i], i] }.first(checked)
+      Ranking.new(labels: order.map { |i| labels[i] }, scores: order.map { |i| scores[i] },
+                  passthrough: false, n: items.length, k: checked)
+    end
+
+    def criteria_items(criteria)
+      items = case criteria
+              when Hash then criteria.to_a
+              when Array then criteria.map { |label| [label, nil] }
+              else raise TypeError, "choice criteria must be a Hash or an Array, got #{criteria.class}"
+              end
+      raise ArgumentError, "choice criteria must contain at least one option" if items.empty?
+
+      duplicate = items.map(&:first).tally.find { |_label, count| count > 1 }
+      raise ArgumentError, "choice criteria label #{duplicate.first.inspect} is duplicated" if duplicate
+
+      items
+    end
+
+    def option_texts(items)
+      Common.render_options({ t: "choice", ins: "", crit: items.to_h }).map(&:to_s)
+    end
+
+    def query_text(state, instructions)
+      body = Common.serialize_state(state)
+      return body if instructions.nil? || instructions == ""
+
+      instructions = PyJSON.dumps(instructions) unless instructions.is_a?(String)
+      "#{instructions}\n#{body}"
+    end
+
+    def narrow(definition, labels)
+      criteria = Util.get(definition, "criteria")
+      kept = criteria.is_a?(Hash) ? labels.to_h { |label| [label, criteria[label]] } : labels.dup
+      Util.put(definition.dup, "criteria", kept)
+    end
+
+    # Whatever `embed_fn` returns, as rows of Float with non-finite values zeroed.
+    def embeddings(embed_fn, texts)
+      raise TypeError, "embed_fn must respond to #call" unless embed_fn.respond_to?(:call)
+
+      rows = rowify(embed_fn.call(texts.dup))
+      unless rows.is_a?(Array) && rows.length == texts.length &&
+             rows.all? { |row| row.is_a?(Array) && !row.empty? } && rows.map(&:length).uniq.length <= 1
+        raise ArgumentError, "embed_fn must return #{texts.length} vectors of equal width, got #{shape(rows)}"
+      end
+
+      rows.map { |row| row.map { |value| finite(value) } }
+    end
+
+    def rowify(raw)
+      return raw if raw.is_a?(Array)
+      return raw.to_a if raw.respond_to?(:to_a)
+
+      raw
+    end
+
+    def shape(rows)
+      return rows.class.to_s unless rows.is_a?(Array)
+      return "#{rows.length} values" unless rows.first.is_a?(Array)
+
+      "#{rows.length} x #{rows.map { |row| row.is_a?(Array) ? row.length : 1 }.uniq.join('|')}"
+    end
+
+    def finite(value)
+      number = Float(value)
+      number.finite? ? number : 0.0
+    rescue ArgumentError, TypeError
+      0.0
+    end
+
+    # Cosine similarity, clipped to [-1, 1] so rounding never reports an impossible score.
+    def cosine(query, documents)
+      norm = Math.sqrt(query.sum { |value| value * value })
+      return Array.new(documents.length, 0.0) if norm.zero? || documents.empty?
+
+      documents.map do |document|
+        length = Math.sqrt(document.sum { |value| value * value })
+        next 0.0 unless (length * norm).positive?
+
+        (document.each_with_index.sum { |value, i| value * query[i] } / (length * norm)).clamp(-1.0, 1.0)
+      end
+    end
+
+    def predict_with(agent, state, questions, **)
+      return agent.predict(state, questions, **) if agent.respond_to?(:predict)
+      return agent.system_one(state, questions, **) if agent.respond_to?(:system_one)
+
+      raise TypeError, "agent must respond to #predict or #system_one"
+    end
+
+    def attach(result, meta)
+      return result.with_shortlist(meta) if result.respond_to?(:with_shortlist)
+      return result.merge("shortlist" => meta) if result.is_a?(Hash)
+
+      raise TypeError, "predict must return a Laya::Result or a Hash, got #{result.class}"
+    end
+  end
+end