ask/answer.ts

import Anthropic from "@anthropic-ai/sdk";
import { fileURLToPath } from "node:url";
import { LAWS, STATE_NAME, type Law } from "./corpus.js";
import { ANSWER_MODEL, requireApiKey, textOf } from "./setup.js";

export const SCOPES = ["road-law", "other"] as const;
export type Scope = (typeof SCOPES)[number];

// The rules. Each is a sentence that was added after a measured failure, and
// each is held by name in a test, so a rewrite that lost one fails before
// anybody reads an answer.
export const RULES = {
  sourced:
    "Every statement comes from the laws given with the question. Nothing from memory about the law, this state or any other state. Put the ids of the laws you used in \"sources\". If no law given touches the question, leave \"answer\" and \"sources\" empty: the app says so itself.",
  figures:
    "Never write a number the laws given do not state. Write figures in digits.",
  register:
    "Write for a sixteen-year-old who asked a friend who knows: plain words, \"you\", each figure once, at most two sentences of under about 20 words each, and no law's name, section or page in the sentences; the app shows those underneath.",
  kinds:
    "A law marked as the court's fine schedule is what a driver actually pays for breaking the law listed above it. A law marked as the state's own page is the licensing agency's word on its own process: prefer it for process, and never for a rule of the road when the manual or the state's law given states the rule.",
  scope:
    "Scope. A rule of the road, a sign, a speed, a fine, an age, a document or a requirement of this state is \"road-law\". Anything else is \"other\", with \"answer\" and \"sources\" empty.",
} as const;

export function rulesText(): string {
  return [
    `You answer a learner driver's question about ${STATE_NAME} road law, from the laws given with the question and nothing else.`,
    "",
    ...Object.values(RULES).map((rule, index) => `${index + 1}. ${rule}`),
  ].join("\n");
}

export const ANSWER_SCHEMA = {
  type: "object",
  properties: {
    scope: { type: "string", enum: [...SCOPES] },
    answer: { type: "string" },
    sources: { type: "array", items: { type: "string" } },
  },
  required: ["scope", "answer", "sources"],
  additionalProperties: false,
} as const;

export type Reply =
  | { kind: "reply"; scope: Scope; answer: string; sources: string[] }
  // Why nothing came back. Written to the log, so a run of failed calls reads
  // differently from a run of refusals.
  | { kind: "failed"; why: "unreachable" | "unparsed" | "refused" };

/** One earlier turn as the learner saw it: what was shown, not what the model first wrote. */
export interface Turn {
  question: string;
  answer: string;
  sources: string[];
}

/** The reference alone: the page, or the section, without the quoted words after it. */
export function referenceOf(citation: string): string {
  const quoted = citation.indexOf(" — ");
  return quoted < 0 ? citation : citation.slice(0, quoted);
}

// The question as the model receives it: the learner's words, then the laws,
// each labelled by what it is. In the user turn, not the system prompt: the
// rules stay the same bytes on every request, and the laws change.
export function withLaws(question: string, laws: Law[]): string {
  if (laws.length === 0) {
    return `${question}\n\n(No law in the learner's state's manual or code was found for this question.)`;
  }
  const listed = laws.map((law) => {
    const reference = referenceOf(law.citation);
    switch (law.kind) {
      case "statute":
        return `[${law.id}] (the state's law, not the manual) ${law.law} (${reference})`;
      case "penalty":
        return `[${law.id}] (the court's fine schedule) ${law.law} (${reference})`;
      case "process":
        return `[${law.id}] (the state's own page) ${law.law} (${reference})`;
      default:
        return `[${law.id}] ${law.law} (manual: ${reference})`;
    }
  });
  return `${question}\n\nLaws from the learner's state retrieved for this question:\n${listed.join("\n")}`;
}

// The request as the SDK takes it, in a function of its own so that the
// measurement in Chapter 11 sends the very same request and cannot drift
// from it.
export function requestFor(
  question: string,
  laws: Law[],
  history: Turn[],
): Anthropic.MessageCreateParamsNonStreaming {
  return {
    model: ANSWER_MODEL,
    max_tokens: 1024,
    system: rulesText(),
    output_config: {
      // Two sentences from a page of law is not a reasoning task.
      effort: "low",
      format: { type: "json_schema", schema: ANSWER_SCHEMA },
    },
    messages: [
      // Replayed as logged: the assistant's turn is what the learner was
      // shown, so a refused turn is a refused turn in the thread too.
      ...history.flatMap((turn): Anthropic.MessageParam[] => [
        { role: "user", content: turn.question },
        {
          role: "assistant",
          content: JSON.stringify({ scope: "road-law", answer: turn.answer, sources: turn.sources }),
        },
      ]),
      { role: "user", content: withLaws(question, laws) },
    ],
  };
}

// What the model's message is read as: a reply in the fixed shape, or why not.
export function replyFrom(response: Anthropic.Message): Reply {
  if (response.stop_reason === "refusal") return { kind: "failed", why: "refused" };
  let parsed: { scope?: unknown; answer?: unknown; sources?: unknown };
  try {
    parsed = JSON.parse(textOf(response)) as typeof parsed;
  } catch {
    return { kind: "failed", why: "unparsed" };
  }
  if (
    typeof parsed.scope !== "string" ||
    !(SCOPES as readonly string[]).includes(parsed.scope) ||
    typeof parsed.answer !== "string" ||
    !Array.isArray(parsed.sources)
  ) {
    return { kind: "failed", why: "unparsed" };
  }
  return {
    kind: "reply",
    scope: parsed.scope as Scope,
    answer: parsed.answer,
    sources: parsed.sources.filter((id): id is string => typeof id === "string"),
  };
}

export type Ask = (question: string, laws: Law[], history: Turn[]) => Promise<Reply>;

// The asker the pipeline uses. A thrown error is a failure with a name, not
// an exception the caller has to know about.
export function askerFor(client: Pick<Anthropic, "messages">): Ask {
  return async (question, laws, history) => {
    try {
      return replyFrom(await client.messages.create(requestFor(question, laws, history)));
    } catch (error) {
      console.error("could not reach the model:", error instanceof Error ? error.message : error);
      return { kind: "failed", why: "unreachable" };
    }
  };
}

async function main(): Promise<void> {
  requireApiKey();
  const question = process.argv.slice(2).join(" ");
  if (question.length === 0) {
    console.error('usage: tsx ask/answer.ts "<question>"');
    process.exit(1);
  }
  const ask = askerFor(new Anthropic());
  // Every law in the corpus, for now. Chapter 4 is where choosing starts.
  console.log(JSON.stringify(await ask(question, LAWS, []), null, 2));
}

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