ask/pipeline.ts

import type { Ask, Turn } from "./answer.js";
import { checkAnswer, pointerFor } from "./check.js";
import type { Law } from "./corpus.js";
import { retrieve, type Search } from "./fuse.js";
import { chooseLaws, poolSize, type Rerank } from "./rerank.js";
import type { Rephrase } from "./rephrase.js";
import { RETRIEVED_LAWS } from "./retrieve.js";
import type { PenaltyRow } from "./schedule.js";
import { pagesLast, wholeSections, withPenalties } from "./sections.js";
import { ANSWER_MODEL } from "./setup.js";
import type { NearLaws } from "./vectors.js";

// One turn, from a question to what the learner is shown and what the log
// records. Everything between the request and the screen is here, so a test
// whose model answers from a script runs the same retrieval and the same
// checks as the deployed pipeline rather than a copy of them.

export interface Pipeline {
  ask: Ask;
  /** The state's laws: the one file a question is answered from. */
  laws: Law[];
  /** The ids nearest the question in meaning. Absent, or throwing, leaves words alone. */
  near?: NearLaws;
  /** The question again in the law's words. Absent, or throwing, retrieval runs on the question alone. */
  rephrase?: Rephrase;
  /** The reranker. Absent, or throwing, the model gets retrieval's first eight. */
  rerank?: Rerank;
  /** The court's fine schedule. Absent, the model is given the sections without it. */
  penalties?: PenaltyRow[];
}

export interface Citation {
  id: string;
  location: string;
  kind?: Law["kind"];
  asOf?: string;
}

export type Body =
  | { id: "answer"; answer: string; citations: Citation[] }
  | { id: "no-match"; pointer: string };

/** What the log records about a turn: what was shown, and why, if nothing was. */
export interface Outcome {
  answer: string;
  sources: string[];
  model: string;
}

function note(what: string, error: unknown): void {
  console.error(`${what}:`, error instanceof Error ? error.message : error);
}

export async function answerQuestion(
  question: string,
  history: Turn[],
  pipeline: Pipeline,
): Promise<{ body: Body; outcome: Outcome }> {
  const { laws } = pipeline;

  let near: string[] = [];
  if (pipeline.near !== undefined) {
    try {
      near = await pipeline.near(question);
    } catch (error) {
      note("the vector index did not answer", error);
    }
  }

  let searches: Search[] = [];
  if (pipeline.rephrase !== undefined) {
    try {
      const worded = await pipeline.rephrase(question);
      searches = await Promise.all(
        worded.map(async (search) => ({
          search,
          near: await (pipeline.near?.(search) ?? Promise.resolve([])).catch(() => []),
        })),
      );
    } catch (error) {
      note("the rephraser did not answer", error);
    }
  }

  const fallback = retrieve(laws, question, RETRIEVED_LAWS, near, searches);
  const pool = pipeline.rerank === undefined ? [] : retrieve(laws, question, poolSize(searches), near, searches);
  const chosen = await chooseLaws(pool, fallback, question, pipeline.rerank);
  const sectioned = wholeSections(chosen, laws);
  const given = pagesLast(
    pipeline.penalties === undefined ? sectioned : withPenalties(sectioned, pipeline.penalties, chosen),
  );

  const reply = await pipeline.ask(question, given, history);

  if (reply.kind === "failed") {
    return {
      body: { id: "no-match", pointer: pointerFor("failed") },
      outcome: { answer: "", sources: [], model: reply.why },
    };
  }
  if (reply.scope !== "road-law") {
    return {
      body: { id: "no-match", pointer: pointerFor("other") },
      outcome: { answer: "", sources: reply.sources, model: `scope:${reply.scope}` },
    };
  }

  const checked = checkAnswer(reply, given, question);
  if (!checked.ok) {
    return {
      body: { id: "no-match", pointer: pointerFor(checked.reason) },
      outcome: { answer: "", sources: reply.sources, model: `road-law:${checked.reason}` },
    };
  }
  return {
    body: {
      id: "answer",
      answer: checked.answer,
      citations: checked.laws.map((law) => ({
        id: law.id,
        location: law.citation,
        ...(law.kind === undefined ? {} : { kind: law.kind }),
        ...(law.asOf === undefined ? {} : { asOf: law.asOf }),
      })),
    },
    outcome: { answer: checked.answer, sources: checked.laws.map((law) => law.id), model: ANSWER_MODEL },
  };
}