aireview 2.1.0 → 2.2.0

data/lib/aireview/jev_critic.rb

@@ -0,0 +1,387 @@
+# frozen_string_literal: true
+require 'json'
+require 'logger'
+require 'yaml'
+require_relative 'utils'
+require_relative 'candidate_checker'
+require_relative 'jev_client'
+
+module Aireview
+  # Jev as a critic. Every candidate becomes a few questions
+  # (prompts/jev_questions.yml) over one state: the MR and Jira sections, the
+  # candidates with the hunks they point at, the diff. The answers turn into
+  # keep, reject or unverifiable by thresholds, then duplicates among the
+  # kept ones are dropped. Jev cannot rewrite a finding, so there is no
+  # refinement. Used as the critique engine (JevStage) and in shadow mode.
+  class JevCritic
+    QUESTIONS = YAML.safe_load_file(File.expand_path('prompts/jev_questions.yml', __dir__)).freeze
+    # The questions whose answers decide; severity only goes to the log. Their
+    # templates go into the review key whole, not as asked about a stub
+    # candidate: a single stub gets no duplicate_of question.
+    DECISION_QUESTIONS = %w[real_issue enough_context version_claim duplicate_of].freeze
+    # Jev limits (docs.typesafe.ai/models): the state plus the longest
+    # question, and the state plus all questions together. Questions count
+    # in both, so the state budget is what they leave.
+    STATE_WITH_LONGEST_QUESTION_TOKENS = 32_000
+    REQUEST_TOKENS = 64_000
+    # There is no Jev tokenizer here: characters are converted pessimistically
+    # (code and Cyrillic take more tokens than English prose) with a margin.
+    # Every request logs the real input_tokens to compare against.
+    CHARS_PER_TOKEN = 2.5
+    SAFETY = 0.9
+    CANDIDATE_FIELDS = %w[id file line quoted_code category severity problem why suggestion note].freeze
+    TRUNCATED = '[truncated to fit the Jev request]'
+    NOT_SHOWN = '(the file is not shown in the diff)'
+    TASK = 'Second pass of a merge request review: check the candidate findings of the first pass ' \
+           'against the diff and the requirements.'
+    SEVERITY_RANK = {'critical' => 0, 'major' => 1, 'minor' => 2}.freeze
+    # The docs only say that a 422 body names the offending field; a size
+    # rejection is recognized by an explicit phrase about going over a limit.
+    # Bare words like "context" are not enough: they occur in question names
+    # (enough_context_C1). Any other 422 is a malformed request, and asking
+    # again would not help.
+    SIZE_ERROR = /
+      \bexceed(?:s|ed|ing)?\b |
+      \btoo\s+(?:long|large|many\s+tokens)\b |
+      \bmaximum\s+(?:context|length|tokens?)\b |
+      \b(?:token|context|length)\s+limit\b
+    /ix
+
+    # decision — :keep, :reject or :unverifiable; answers — the numbers
+    # behind it, for the log.
+    Assessment = Struct.new(:id, :decision, :reason, :answers, keyword_init: true) do
+      # Rounded for reading; the exact numbers stay in answers.
+      def numbers
+        answers.map do |name, value|
+          value.is_a?(Hash) ? "#{name}=#{value[:choice]}/#{round(value[:confidence])}" : "#{name}=#{round(value)}"
+        end.join(' ')
+      end
+
+      private
+
+      def round(value)
+        value.is_a?(Numeric) ? value.round(2) : value
+      end
+    end
+    # answers — every answer by question key, for dropping duplicates once
+    # the verdicts of Jev and of the LLM are merged.
+    Result = Struct.new(:assessments, :answers, :requests, :model, keyword_init: true)
+    Request = Struct.new(:candidates, :state, :questions, keyword_init: true)
+    # What the requests of one assessment gathered.
+    Asked = Struct.new(:answers, :unfit, :requests, :model, keyword_init: true)
+
+    # scrub — the secret scrubber of the context: the candidates are LLM text
+    # and are scrubbed like the candidates JSON of the Critique prompt.
+    def initialize(client:, thresholds:, review_instructions: nil, scrub: ->(text) { text },
+                   logger: Logger.new($stderr))
+      @client = client
+      @thresholds = thresholds
+      @scrub = scrub
+      instructions = Aireview::Utils.presence(review_instructions.to_s.strip)
+      @review_instructions = instructions && scrub.call(instructions)
+      @logger = logger
+    end
+
+    # Kept candidates in generate order: an edge comes from a
+    # duplicate_of_<id> answer that names another kept id with enough
+    # confidence. In a group of duplicates the most severe one stays, on a
+    # tie the one generate listed first. Returns {dropped id => kept id}.
+    # The client and the critic as the config sets them up; scrub — the
+    # secret scrubber of the context.
+    def self.build(config:, scrub:, logger:)
+      new(client: JevClient.new(config: config, logger: logger), thresholds: config.jev_thresholds,
+          review_instructions: config.review_instructions, scrub: scrub, logger: logger)
+    end
+
+    def self.decision_templates
+      QUESTIONS.slice(*DECISION_QUESTIONS)
+    end
+
+    def self.duplicates(kept, answers, threshold:)
+      ids = kept.map { |candidate| field(candidate, 'id') }
+      rank = kept.to_h do |candidate|
+        [field(candidate, 'id'), SEVERITY_RANK.fetch(field(candidate, 'severity'), SEVERITY_RANK.size)]
+      end
+      duplicate_groups(ids, answers, threshold).each_with_object({}) do |group, dropped|
+        winner = group.min_by { |id| [rank[id], ids.index(id)] }
+        (group - [winner]).each { |id| dropped[id] = winner }
+      end
+    end
+
+    def self.duplicate_groups(ids, answers, threshold)
+      group_of = ids.to_h { |id| [id, [id]] }
+      ids.each do |id|
+        other = duplicate_answer(id, ids, answers, threshold)
+        join_groups(group_of, id, other) if other
+      end
+      group_of.values.uniq(&:object_id).select { |group| group.size > 1 }
+    end
+
+    def self.join_groups(group_of, id, other)
+      return if group_of[id].equal?(group_of[other])
+
+      merged = group_of[id] + group_of[other]
+      merged.each { |member| group_of[member] = merged }
+    end
+
+    # The other kept id a duplicate_of answer names with enough confidence.
+    def self.duplicate_answer(id, ids, answers, threshold)
+      answer = answers["duplicate_of_#{id}"]
+      return nil unless answer.is_a?(Hash) && answer['confidence'].to_f >= threshold
+
+      other = answer['choice']
+      other if other != id && ids.include?(other)
+    end
+
+    def self.field(candidate, key)
+      (candidate[key] || candidate[key.to_sym]).to_s.strip
+    end
+
+    # dedupe — false when the caller merges the verdicts with the LLM ones
+    # first and drops duplicates over the whole set (JevStage).
+    def assess(context:, candidates:, dedupe: true)
+      candidates = candidates.map { |candidate| normalize(candidate) }
+      sections = diff_sections(context.diff_text)
+      asked = Asked.new(answers: {}, unfit: [], requests: 0, model: nil)
+      plan_requests(context, candidates, sections).each { |request| ask(request, asked, context, sections) }
+      assessments = candidates.map { |candidate| decide(candidate, asked.answers, asked.unfit) }
+      assessments = drop_duplicates(assessments, candidates, asked.answers) if dedupe
+      Result.new(assessments: assessments, answers: asked.answers, requests: asked.requests, model: asked.model)
+    end
+
+    # The first request as it would be sent, for --dry-run.
+    def preview(context:, candidates:)
+      candidates = candidates.map { |candidate| normalize(candidate) }
+      plan_requests(context, candidates, diff_sections(context.diff_text)).first
+    end
+
+    private
+
+    def normalize(candidate)
+      CANDIDATE_FIELDS.to_h do |field|
+        value = candidate[field] || candidate[field.to_sym]
+        [field, value.is_a?(String) ? @scrub.call(value) : value]
+      end.compact.merge('id' => self.class.field(candidate, 'id'))
+    end
+
+    # One request for all candidates when it fits, otherwise one per
+    # candidate (then without duplicate_of: a request sees one candidate).
+    # A candidate that does not fit even alone is unverifiable.
+    def plan_requests(context, candidates, sections)
+      shared = build_request(context, candidates, sections)
+      return [shared] if shared.state || candidates.size == 1
+
+      @logger.info('Jev: the candidates do not fit one request, asking about each separately')
+      candidates.map { |candidate| build_request(context, [candidate], sections) }
+    end
+
+    def build_request(context, candidates, sections)
+      ids = candidates.map { |candidate| candidate['id'] }
+      questions = candidates.map { |candidate| questions_for(candidate['id'], ids) }.reduce({}, :merge)
+      state = {
+        'task' => TASK,
+        'requirements' => context.sections,
+        'project_instructions' => @review_instructions,
+        'context_truncated' => coverage_notes(context.coverage),
+        'candidates' => candidates.to_h { |candidate| [candidate['id'], with_evidence(candidate, sections)] },
+        'diff' => context.diff_text
+      }.compact
+      Request.new(candidates: candidates, state: fit(state, state_budget(questions)), questions: questions)
+    end
+
+    def questions_for(id, ids)
+      questions = %w[real_issue enough_context version_claim severity].to_h do |name|
+        ["#{name}_#{id}", question(QUESTIONS.fetch(name), id)]
+      end
+      others = ids - [id]
+      questions["duplicate_of_#{id}"] = duplicate_question(id, others) unless others.empty?
+      questions
+    end
+
+    def question(template, id)
+      {
+        'type' => template.fetch('type'),
+        'instructions' => format(template.fetch('instructions'), id: id),
+        'criteria' => template.fetch('criteria').transform_values { |text| format(text, id: id) }
+      }
+    end
+
+    def duplicate_question(id, others)
+      template = QUESTIONS.fetch('duplicate_of')
+      criteria = others.to_h { |other| [other, format(template.fetch('other'), other: other)] }
+      {
+        'type' => template.fetch('type'),
+        'instructions' => format(template.fetch('instructions'), id: id),
+        'criteria' => criteria.merge('none' => template.fetch('none'))
+      }
+    end
+
+    def state_budget(questions)
+      sizes = questions.values.map { |question| tokens(JSON.generate(question).length) }
+      room = [STATE_WITH_LONGEST_QUESTION_TOKENS - sizes.max, REQUEST_TOKENS - sizes.sum].min
+      (room * SAFETY * CHARS_PER_TOKEN).floor
+    end
+
+    def tokens(chars)
+      (chars / CHARS_PER_TOKEN).ceil
+    end
+
+    # The diff goes first, then the MR and Jira sections; the candidates and
+    # their evidence are never cut. nil — does not fit even then.
+    def fit(state, budget)
+      %w[diff requirements].each do |key|
+        break if size(state) <= budget
+
+        state = shrink(state, key, budget) if state[key]
+      end
+      size(state) <= budget ? state : nil
+    end
+
+    # Every character of the text takes at least one character of JSON, so
+    # cutting the overflow in characters is enough; the marker costs its
+    # length plus the escaped newline.
+    def shrink(state, key, budget)
+      text = state[key].to_s
+      keep = text.length - (size(state) - budget) - TRUNCATED.length - 2
+      state.merge(key => keep.positive? ? "#{text[0, keep]}\n#{TRUNCATED}" : TRUNCATED)
+    end
+
+    def size(state)
+      JSON.generate(state).length
+    end
+
+    def coverage_notes(coverage)
+      return nil if coverage.nil? || coverage.complete?
+
+      notes = coverage.truncated_sections.map { |label| "#{label} truncated" }
+      notes << "files not shown: #{coverage.files_not_shown.join(', ')}" unless coverage.files_not_shown.empty?
+      notes += coverage.files_partial.map { |file| "#{file[:path]}: #{file[:shown]} of #{file[:total]} hunks shown" }
+      notes << "diff not available: #{coverage.files_unavailable.join(', ')}" unless coverage.files_unavailable.empty?
+      notes
+    end
+
+    # The hunk the line falls into; the whole shown diff of the file when
+    # the line is unknown or the anchoring check left a note.
+    def with_evidence(candidate, sections)
+      section = sections[candidate['file'].to_s] || sections[candidate['file'].to_s.sub(%r{\A(?:\./|[ab]/)}, '')]
+      evidence = if section.nil?
+                   NOT_SHOWN
+                 elsif candidate['line'].is_a?(Integer) && !candidate['note']
+                   hunk = section[:hunks].find { |item| item[:range].cover?(candidate['line']) }
+                   hunk ? hunk[:text] : section[:text]
+                 else
+                   section[:text]
+                 end
+      candidate.merge('evidence' => evidence)
+    end
+
+    # The shown diff by file: its whole text and every hunk with the range of
+    # new-file lines it covers (the headers as CandidateChecker reads them).
+    def diff_sections(diff_text)
+      sections = {}
+      section = nil
+      diff_text.to_s.each_line do |line|
+        if (header = line.match(CandidateChecker::FILE_HEADER))
+          section = {text: +'', hunks: []}
+          header.captures.map(&:strip).each { |path| sections[path] = section }
+        elsif section && (hunk_header = line.match(CandidateChecker::HUNK_HEADER))
+          section[:hunks] << new_hunk(hunk_header)
+        end
+        add_line(section, line) if section
+      end
+      sections
+    end
+
+    def new_hunk(header)
+      start = header[1].to_i
+      length = header[2] ? header[2].to_i : 1
+      {range: start..(start + length - 1), text: +''}
+    end
+
+    def add_line(section, line)
+      section[:text] << line
+      section[:hunks].last[:text] << line unless section[:hunks].empty?
+    end
+
+    # The estimate goes to the log next to the input_tokens Jev reports, to
+    # check CHARS_PER_TOKEN against real requests. When Jev rejects a request
+    # as too large after all (the estimate missed), its candidates are asked
+    # about one by one; one that is too large alone is unverifiable.
+    def ask(request, asked, context, sections)
+      ids = request.candidates.map { |candidate| candidate['id'] }
+      return asked.unfit.concat(ids) unless request.state
+
+      chars = size(request.state) + request.questions.values.sum { |question| JSON.generate(question).length }
+      @logger.info("Jev request estimated at ~#{tokens(chars)} tokens (#{chars} chars of state and questions)")
+      asked.requests += 1
+      result = @client.evaluate(state: request.state, questions: request.questions)
+      asked.answers.merge!(result.answers)
+      asked.model = result.model
+    rescue JevError => e
+      raise unless e.status == 422 && e.message.match?(SIZE_ERROR)
+
+      ask_separately(request, asked, context, sections, e)
+    end
+
+    def ask_separately(request, asked, context, sections, error)
+      ids = request.candidates.map { |candidate| candidate['id'] }
+      if ids.size == 1
+        @logger.warn("Jev rejected the request about #{ids.first} as too large, " \
+                     "it stays unverifiable: #{error.message}")
+        return asked.unfit.concat(ids)
+      end
+
+      @logger.warn("Jev rejected the request as too large, asking about each candidate separately: #{error.message}")
+      request.candidates.each do |candidate|
+        ask(build_request(context, [candidate], sections), asked, context, sections)
+      end
+    end
+
+    def decide(candidate, answers, unfit)
+      id = candidate['id']
+      numbers = numbers(id, answers)
+      assessment = ->(decision, reason) { Assessment.new(id: id, decision: decision, reason: reason, answers: numbers) }
+      if unfit.include?(id)
+        return assessment.call(:unverifiable, 'the candidate and its evidence do not fit a Jev request')
+      end
+      if numbers[:version_claim] >= @thresholds[:version_claim]
+        return assessment.call(:reject, 'claims that a version does not exist')
+      end
+      if numbers[:enough_context] < @thresholds[:enough_context]
+        return assessment.call(:unverifiable, 'not enough context in the request')
+      end
+
+      if numbers[:real_issue] >= @thresholds[:keep_above]
+        assessment.call(:keep, 'real issue')
+      else
+        assessment.call(:reject, 'not a confirmed issue')
+      end
+    end
+
+    def numbers(id, answers)
+      noul = ->(name) { answers.dig("#{name}_#{id}", 'noul') }
+      choice = lambda do |name|
+        answer = answers["#{name}_#{id}"]
+        answer && {choice: answer['choice'], confidence: answer['confidence']}
+      end
+      {
+        real_issue: noul.call('real_issue'), enough_context: noul.call('enough_context'),
+        version_claim: noul.call('version_claim'), duplicate_of: choice.call('duplicate_of'),
+        severity: choice.call('severity')
+      }.compact
+    end
+
+    def drop_duplicates(assessments, candidates, answers)
+      kept_ids = assessments.select { |assessment| assessment.decision == :keep }.map(&:id)
+      kept = candidates.select { |candidate| kept_ids.include?(candidate['id']) }
+      dropped = self.class.duplicates(kept, answers, threshold: @thresholds[:duplicate])
+      assessments.map do |assessment|
+        winner = dropped[assessment.id]
+        next assessment unless winner
+
+        Assessment.new(id: assessment.id, decision: :reject, reason: "duplicate of #{winner}",
+                       answers: assessment.answers)
+      end
+    end
+  end
+end

data/lib/aireview/jev_shadow.rb

@@ -0,0 +1,82 @@
+# frozen_string_literal: true
+require 'json'
+require 'logger'
+require_relative 'errors'
+require_relative 'utils'
+require_relative 'jev_critic'
+
+module Aireview
+  # llm.jev.shadow: after the LLM Critique the same candidates go to Jev, and
+  # its decisions are logged next to the Critique verdicts. The review result
+  # and the review key do not depend on it, and a Jev failure is a warning:
+  # this is data for choosing thresholds, not a check.
+  class JevShadow
+    def initialize(config:, scrub:, logger: Logger.new($stderr), critic: nil)
+      @config = config
+      @scrub = scrub
+      @logger = logger
+      @critic = critic
+    end
+
+    # accepted — what the LLM Critique kept; only the ids are compared.
+    def run(context:, candidates:, accepted:)
+      critic = self.critic
+      return @logger.warn('Jev shadow skipped: JEV_API_KEY is not set') unless critic
+
+      started_at = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+      result = critic.assess(context: context, candidates: candidates)
+      seconds = Process.clock_gettime(Process::CLOCK_MONOTONIC) - started_at
+      log(result, accepted.map { |candidate| (candidate['id'] || candidate[:id]).to_s }, seconds)
+    rescue JevError => e
+      @logger.warn("Jev shadow failed, the review is not affected: #{e.message}")
+    # The shadow is an experiment running in real review jobs: a bug in it
+    # must not cost a review either, but it must stay visible as a bug.
+    rescue StandardError => e
+      @logger.warn("Jev shadow crashed, the review is not affected: #{e.class}: #{e.message} " \
+                   "(#{e.backtrace&.first})")
+    end
+
+    private
+
+    def critic
+      return @critic if @critic
+      return nil if Aireview::Utils.blank?(@config.jev_api_key)
+
+      @critic = JevCritic.build(config: @config, scrub: @scrub, logger: @logger)
+    end
+
+    def log(result, critique_kept, seconds)
+      result.assessments.each do |assessment|
+        critique = critique_kept.include?(assessment.id) ? 'keep' : 'reject'
+        @logger.info("Jev shadow #{assessment.id}: #{assessment.decision} (#{assessment.reason}; " \
+                     "#{assessment.numbers}), critique: #{critique}")
+      end
+      @logger.info("Jev shadow: #{summary(result.assessments, critique_kept)} " \
+                   "(model=#{result.model}, requests=#{result.requests}, #{format('%.1fs', seconds)})")
+      @logger.info("Jev shadow data: #{JSON.generate(data(result, critique_kept))}")
+    end
+
+    # The lines above round for reading; thresholds are chosen from this one,
+    # which keeps the numbers exactly as Jev returned them.
+    def data(result, critique_kept)
+      {
+        model: result.model,
+        candidates: result.assessments.map do |assessment|
+          {id: assessment.id, decision: assessment.decision, reason: assessment.reason,
+           critique: critique_kept.include?(assessment.id) ? 'keep' : 'reject', **assessment.answers}
+        end
+      }
+    end
+
+    def summary(assessments, critique_kept)
+      decided = assessments.reject { |assessment| assessment.decision == :unverifiable }
+      agreed = decided.count do |assessment|
+        (assessment.decision == :keep) == critique_kept.include?(assessment.id)
+      end
+      counts = %i[keep reject unverifiable].map do |decision|
+        "#{decision} #{assessments.count { |assessment| assessment.decision == decision }}"
+      end
+      "agrees with critique on #{agreed} of #{decided.size} decided candidate(s); #{counts.join(', ')}"
+    end
+  end
+end

data/lib/aireview/jev_stage.rb

@@ -0,0 +1,97 @@
+# frozen_string_literal: true
+require 'logger'
+require_relative 'errors'
+require_relative 'jev_critic'
+
+module Aireview
+  # The Critique stage with llm.critique.engine: jev. Jev decides keep or
+  # reject; the candidates it cannot judge (not enough context, too large)
+  # and, when Jev fails, all of them go to the LLM Critique with
+  # llm.jev.fallback: model, or are rejected / fail the run with fail.
+  # Duplicates are dropped once, over the merged verdicts of Jev and the LLM:
+  # neither side sees what the other kept.
+  class JevStage
+    # The report line: :jev (Jev decided everything), :jev_partial (some
+    # went to the LLM), :jev_failed (the LLM decided everything).
+    Outcome = Struct.new(:accepted, :note, keyword_init: true)
+
+    def initialize(config:, critic:, logger: Logger.new($stderr))
+      @config = config
+      @critic = critic
+      @logger = logger
+    end
+
+    # llm_critique — the LLM Critique of the pipeline for a subset of the
+    # candidates; returns the kept ones, refined.
+    def run(context:, candidates:, llm_critique:)
+      @logger.info("Pipeline critique pass started (engine=jev, model=#{@config.jev_model})")
+      begin
+        result = @critic.assess(context: context, candidates: candidates, dedupe: false)
+      rescue JevError => e
+        return fall_back(e, candidates, llm_critique)
+      end
+      result.assessments.each { |assessment| log(assessment) }
+
+      unverifiable = pick(candidates, result, :unverifiable)
+      accepted = pick(candidates, result, :keep) + judge_unverifiable(unverifiable, llm_critique)
+      accepted = drop_duplicates(in_generate_order(accepted, candidates), result.answers)
+      @logger.info("Pipeline critique pass completed with #{accepted.size} kept candidate(s) (engine=jev, " \
+                   "model=#{result.model}, requests=#{result.requests})")
+      Outcome.new(accepted: accepted, note: unverifiable.empty? || fail? ? :jev : :jev_partial)
+    end
+
+    private
+
+    def fail?
+      @config.jev_fallback == 'fail'
+    end
+
+    def fall_back(error, candidates, llm_critique)
+      if fail?
+        raise JevError.new("Jev critique failed and llm.jev.fallback is fail: #{error.message}", status: error.status)
+      end
+
+      @logger.warn("Jev critique failed, the LLM critique takes over (llm.jev.fallback: model): #{error.message}")
+      Outcome.new(accepted: llm_critique.call(candidates), note: :jev_failed)
+    end
+
+    def judge_unverifiable(unverifiable, llm_critique)
+      return [] if unverifiable.empty?
+
+      ids = unverifiable.map { |candidate| id_of(candidate) }.join(', ')
+      if fail?
+        @logger.info("Critique reject #{ids}: Jev could not judge them and llm.jev.fallback is fail")
+        return []
+      end
+
+      @logger.info("Pipeline critique: #{ids} go to the LLM critique, Jev could not judge them " \
+                   "(model=#{@config.critique_model})")
+      llm_critique.call(unverifiable)
+    end
+
+    def pick(candidates, result, decision)
+      ids = result.assessments.select { |assessment| assessment.decision == decision }.map(&:id)
+      candidates.select { |candidate| ids.include?(id_of(candidate)) }
+    end
+
+    def in_generate_order(accepted, candidates)
+      order = candidates.map { |candidate| id_of(candidate) }
+      accepted.sort_by { |candidate| order.index(id_of(candidate)) || order.size }
+    end
+
+    def drop_duplicates(accepted, answers)
+      dropped = JevCritic.duplicates(accepted, answers, threshold: @config.jev_thresholds[:duplicate])
+      dropped.each { |id, winner| @logger.info("Critique reject #{id}: duplicate of #{winner} (Jev)") }
+      accepted.reject { |candidate| dropped.key?(id_of(candidate)) }
+    end
+
+    def log(assessment)
+      message = "Critique (jev) #{assessment.decision} #{assessment.id}: #{assessment.reason} (#{assessment.numbers})"
+      assessment.decision == :keep ? @logger.debug(message) : @logger.info(message)
+    end
+
+    def id_of(candidate)
+      (candidate['id'] || candidate[:id]).to_s.strip
+    end
+  end
+end

data/lib/aireview/llm_client.rb

@@ -30,7 +30,7 @@ module Aireview
       stage = prompt.stage.to_s
       model = candidate.model
       @logger.info("LLM #{stage} request started (model=#{model}, temperature=#{prompt.temperature})")
-      chat = build_chat(context: context(stage, candidate.provider, key, key_index), stage: stage,
+      chat = build_chat(context: context(stage, candidate, key, key_index), stage: stage,
                         model: model, provider: candidate.provider)
       chat = configure_reasoning(chat: chat, model: model, provider: candidate.provider)
         .with_temperature(prompt.temperature.to_f)
@@ -99,18 +99,19 @@ module Aireview
       context.chat(model: model, provider: provider.to_sym, assume_model_exists: true)
     end
 
-    # A RubyLLM context per stage, provider and key index: switching the key
-    # is another context, not an edit of the global config.
-    def context(stage, provider, key, key_index)
-      @contexts[[stage, provider, key_index]] ||= build_context(provider.to_s, key)
+    # A RubyLLM context per stage, key source (a provider's API or a server
+    # of your own) and key index: switching the key is another context, not
+    # an edit of the global config.
+    def context(stage, candidate, key, key_index)
+      @contexts[[stage, candidate.key_source, key_index]] ||= build_context(candidate, key)
     end
 
-    def build_context(provider, api_key)
+    def build_context(candidate, api_key)
       RubyLLM.context do |ruby_config|
         configure_http_proxy(ruby_config)
         ruby_config.request_timeout = @config.llm_timeout.to_f
         ruby_config.max_retries = 0
-        configure_provider(ruby_config, provider, api_key)
+        configure_provider(ruby_config, candidate.provider.to_s, api_key, candidate.api_base)
       end
     end
 
@@ -120,27 +121,30 @@ module Aireview
       ruby_config.http_proxy = @config.llm_http_proxy
     end
 
-    def configure_provider(ruby_config, provider, api_key)
+    # api_base — the model's own server; without it the provider's API,
+    # or LLM_API_BASE for the providers that always took it.
+    def configure_provider(ruby_config, provider, api_key, api_base)
       case provider
       when 'gemini', 'openai', 'openrouter'
-        configure_remote_provider(ruby_config, provider, api_key)
+        configure_remote_provider(ruby_config, provider, api_key, api_base || @config.llm_api_base)
       when 'anthropic'
         ruby_config.anthropic_api_key = api_key
+        ruby_config.anthropic_api_base = api_base if api_base
       when 'ollama'
-        ruby_config.ollama_api_base = @config.ollama_api_base
+        ruby_config.ollama_api_base = api_base || @config.ollama_api_base
       else
         raise ConfigError, "Unsupported LLM provider: #{provider.inspect}"
       end
     end
 
     # RubyLLM 2 sends OpenAI requests to the Responses API; a compatible
-    # server behind LLM_API_BASE usually has Chat Completions only.
-    def configure_remote_provider(ruby_config, provider, api_key)
+    # server usually has Chat Completions only.
+    def configure_remote_provider(ruby_config, provider, api_key, api_base)
       ruby_config.public_send("#{provider}_api_key=", api_key)
       ruby_config.openai_protocol = :chat_completions if provider == 'openai'
-      return unless Aireview::Utils.present?(@config.llm_api_base)
+      return unless Aireview::Utils.present?(api_base)
 
-      ruby_config.public_send("#{provider}_api_base=", @config.llm_api_base)
+      ruby_config.public_send("#{provider}_api_base=", api_base)
     end
   end
 end

data/lib/aireview/llm_router.rb

@@ -34,7 +34,7 @@ module Aireview
 
     Slot = Struct.new(:candidate, :index)
     # The key the next model of the same provider continues from.
-    Carry = Struct.new(:provider, :key_index)
+    Carry = Struct.new(:key_source, :key_index)
     Delay = Struct.new(:seconds, :source)
 
     SWITCH_HINT = 'Try again later or switch model via --generate-model/--critique-model.'
@@ -213,7 +213,7 @@ module Aireview
         log_switch(stage, route, visits)
         status, response = try_route(stage, route, visits, &request)
         return [:ok, remember(stage, route, response)] if status == :ok
-        return [:next_model, Carry.new(route.candidate.provider, route.key_index)] if status == :next_model
+        return [:next_model, Carry.new(route.candidate.key_source, route.key_index)] if status == :next_model
       end
       state(slot.candidate).exclude('daily quota exhausted on every key') unless tried
       [:next_model, carry]
@@ -224,7 +224,7 @@ module Aireview
     # is bound to "key + model", a failure on one model does not write the
     # key off for another.
     def candidate_routes(stage, slot, carry)
-      keys = @config.provider_api_keys(slot.candidate.provider)
+      keys = @config.candidate_api_keys(slot.candidate)
       routes = keys.each_with_index.map do |key, key_index|
         Route.new(candidate: slot.candidate, candidate_index: slot.index, key: key,
                   key_index: key_index, key_count: keys.size)
@@ -233,7 +233,7 @@ module Aireview
     end
 
     def start_key(stage, slot, carry)
-      return carry.key_index if carry && carry.provider == slot.candidate.provider
+      return carry.key_index if carry && carry.key_source == slot.candidate.key_source
 
       cursor_candidate, cursor_key = @cursor.fetch(stage, [0, 0])
       cursor_candidate == slot.index ? cursor_key : 0