ask/rerank.ts

import Anthropic from "@anthropic-ai/sdk";
import { fileURLToPath } from "node:url";
import { LAWS, type Law } from "./corpus.js";
import { retrieve } from "./fuse.js";
import { SMALL_MODEL, requireApiKey, textOf } from "./setup.js";
import { sideJobClient } from "./rephrase.js";
import { embeddingText } from "./vectors.js";

// The laws the model is given, chosen by a reranker from a wider pool.
// Retrieval ranks by how a law is worded; a reranker reads the question and
// the laws together and judges which answer it. Absent or failing, the model
// gets what retrieval gave it before the reranker existed.

/** How many retrieved laws the reranker reads. */
export const CANDIDATE_LAWS = 24;

// How many when the question was also searched in the law's words. Fusing a
// search's rankings beside the question's own moves laws out of the pool as
// well as in, so the pool is wider when searches were fused.
export const SEARCHED_CANDIDATE_LAWS = 32;

/** How many it gives the model. */
export const ANSWER_LAWS = 3;

/** How many laws Haiku names, best first. The rest are unranked and keep their place in the pool. */
const HAIKU_RANKS = 5;

export function poolSize(searches: readonly unknown[]): number {
  return searches.length > 0 ? SEARCHED_CANDIDATE_LAWS : CANDIDATE_LAWS;
}

/** A score for each text, in the order given. */
export type Rerank = (question: string, texts: string[]) => Promise<number[]>;

const HAIKU_RULES = `You choose which laws answer a learner driver's question.
You are given the question and a numbered list of laws from one state's driver manual and vehicle code.
Return the numbers of the laws that answer the question, best first, at most ${HAIKU_RANKS}.
A law that states the rule, figure or penalty asked about ranks above one that only names the subject.
Match the situation asked about: an adult's limit is not an under-21 limit, a side window is not a windshield, a first offence is not a second.
Return an empty list if none answers it.`;

const HAIKU_SCHEMA = {
  type: "object",
  properties: { ranked: { type: "array", items: { type: "integer" } } },
  required: ["ranked"],
  additionalProperties: false,
} as const;

// Haiku as a reranker: the laws it names score from HAIKU_RANKS down, in its
// order; every other law is -Infinity and keeps its place. Throws on anything
// it cannot read, so the caller falls back.
export function haikuRerankFor(client: Pick<Anthropic, "messages">): Rerank {
  return async (question, texts) => {
    const listed = texts.map((text, index) => `[${index}] ${text}`).join("\n\n");
    const response = await client.messages.create({
      model: SMALL_MODEL,
      max_tokens: 1024,
      system: HAIKU_RULES,
      output_config: {
        effort: "low",
        format: { type: "json_schema", schema: HAIKU_SCHEMA },
      },
      messages: [{ role: "user", content: `Question: ${question}\n\nLaws:\n\n${listed}` }],
    });
    const text = textOf(response);
    const ranked = (JSON.parse(text) as { ranked?: unknown }).ranked;
    if (
      !Array.isArray(ranked) ||
      ranked.some((id) => !Number.isInteger(id) || (id as number) < 0 || (id as number) >= texts.length)
    ) {
      throw new Error(`the reranker named laws that were not given: ${text}`);
    }
    const scores = texts.map(() => Number.NEGATIVE_INFINITY);
    [...new Set(ranked as number[])]
      .slice(0, HAIKU_RANKS)
      .forEach((id, position) => (scores[id] = HAIKU_RANKS - position));
    return scores;
  };
}

/**
 * The laws the model answers from: the reranker's best ANSWER_LAWS of the
 * pool, or, with no reranker or a failing one, `fallback` -- what retrieval
 * gave before the reranker existed. A row of the state's own page the
 * reranker puts among its first is given beside the three, never in place of
 * one.
 */
export async function chooseLaws(
  pool: Law[],
  fallback: Law[],
  question: string,
  rerank?: Rerank,
): Promise<Law[]> {
  if (rerank === undefined || pool.length === 0) return fallback;
  try {
    const scores = await rerank(question, pool.map(embeddingText));
    if (scores.length !== pool.length) throw new Error(`${scores.length} scores for ${pool.length} laws`);
    const ranked = pool.map((law, index) => ({ law, score: scores[index] })).sort((a, b) => b.score - a.score);
    const page = (law: Law): boolean => law.kind === "process";
    const pages = ranked.slice(0, ANSWER_LAWS).filter(({ law }) => page(law)).map(({ law }) => law);
    const chosen = ranked.filter(({ law }) => !page(law)).slice(0, ANSWER_LAWS).map(({ law }) => law);
    return [...chosen, ...pages];
  } catch (error) {
    console.error("the reranker did not answer:", error instanceof Error ? error.message : error);
    return fallback;
  }
}

async function main(): Promise<void> {
  requireApiKey();
  const question = process.argv.slice(2).join(" ");
  if (question.length === 0) {
    console.error('usage: tsx ask/rerank.ts "<question>"');
    process.exit(1);
  }
  const pool = retrieve(LAWS, question, CANDIDATE_LAWS);
  const fallback = retrieve(LAWS, question);
  for (const law of await chooseLaws(pool, fallback, question, haikuRerankFor(sideJobClient()))) {
    console.log(`${law.id}  ${law.law}`);
  }
}

if (process.argv[1] === fileURLToPath(import.meta.url)) {
  main().catch((error: unknown) => {
    console.error(error instanceof Error ? error.message : String(error));
    process.exit(1);
  });
}