ask/vectors.ts

import { fileURLToPath } from "node:url";
import { LAWS, type Law } from "./corpus.js";

// The meaning half of retrieval. BM25 finds a law that shares the question's
// words; this finds one that shares its meaning -- "can I leave my car right
// by a fireplug", or the same question in Spanish against English law text,
// neither of which shares a word with "within 15 feet of a fire hydrant".
//
// Only ids come back from here. The words of every law shown, and the citation
// under it, are read from the corpus by id and checked exactly as before, so
// the index can rank badly but cannot put text on screen.

/** Multilingual, so a Spanish question lands near the English law that answers it. */
export const EMBEDDING_MODEL = "@cf/baai/bge-m3";

/** How many neighbours a question asks for. */
export const NEAR_LAWS = 8;

// Below this a neighbour is not a match, only the nearest thing there was.
// Set for recall, not precision: a law that does not answer costs a few
// tokens and is declined by the model, with the checker behind it; a law
// that does answer and is cut here is a learner told their manual does not
// say.
export const NEAR_ENOUGH = 0.5;

export type Embed = (texts: string[]) => Promise<number[][]>;

/** What is embedded for a law: its rule and the card's question, which is how a learner would ask it. */
export function embeddingText(law: Pick<Law, "law" | "question">): string {
  return `${law.law} ${law.question}`.trim();
}

// Workers AI over Cloudflare's REST API. Absent credentials mean the stage
// does not exist: the caller gets `undefined` and retrieves by words alone.
export function workersAiEmbedder(): Embed | undefined {
  const account = process.env.CLOUDFLARE_ACCOUNT_ID;
  const token = process.env.CLOUDFLARE_API_TOKEN;
  if (!account || !token) return undefined;
  return async (texts) => {
    const response = await fetch(
      `https://api.cloudflare.com/client/v4/accounts/${account}/ai/run/${EMBEDDING_MODEL}`,
      {
        method: "POST",
        headers: { Authorization: `Bearer ${token}`, "Content-Type": "application/json" },
        body: JSON.stringify({ text: texts }),
      },
    );
    const parsed = (await response.json()) as {
      success: boolean;
      errors: unknown;
      result?: { data: number[][] };
    };
    if (!parsed.success || parsed.result === undefined) {
      throw new Error(`embedding failed: ${response.status} ${JSON.stringify(parsed.errors)}`);
    }
    if (parsed.result.data.length !== texts.length) {
      throw new Error(`asked to embed ${texts.length} texts, given ${parsed.result.data.length} vectors`);
    }
    return parsed.result.data;
  };
}

/** How alike two vectors point: 1 for the same direction, 0 for unrelated. */
export function cosine(a: number[], b: number[]): number {
  let dot = 0;
  let aa = 0;
  let bb = 0;
  for (let i = 0; i < a.length; i += 1) {
    dot += a[i] * b[i];
    aa += a[i] * a[i];
    bb += b[i] * b[i];
  }
  return dot / (Math.sqrt(aa) * Math.sqrt(bb));
}

export interface Index {
  ids: string[];
  vectors: number[][];
}

/** At most this many texts in one embedding request. */
const PER_REQUEST = 100;

export async function buildIndex(laws: Law[], embed: Embed): Promise<Index> {
  const vectors: number[][] = [];
  for (let at = 0; at < laws.length; at += PER_REQUEST) {
    vectors.push(...(await embed(laws.slice(at, at + PER_REQUEST).map(embeddingText))));
  }
  return { ids: laws.map((law) => law.id), vectors };
}

/** The ids nearest a vector, best first, cut below `cutoff`. */
export function nearest(
  index: Index,
  vector: number[],
  topK = NEAR_LAWS,
  cutoff = NEAR_ENOUGH,
): { id: string; score: number }[] {
  return index.ids
    .map((id, i) => ({ id, score: cosine(vector, index.vectors[i]) }))
    .filter((hit) => hit.score >= cutoff)
    .sort((a, b) => b.score - a.score)
    .slice(0, topK);
}

export type NearLaws = (question: string) => Promise<string[]>;

export function nearLawsFor(index: Index, embed: Embed): NearLaws {
  return async (question) => {
    const [vector] = await embed([question]);
    return nearest(index, vector).map((hit) => hit.id);
  };
}

async function main(): Promise<void> {
  const question = process.argv.slice(2).join(" ");
  if (question.length === 0) {
    console.error('usage: tsx ask/vectors.ts "<question>"');
    process.exit(1);
  }
  const embed = workersAiEmbedder();
  if (embed === undefined) {
    console.error("CLOUDFLARE_ACCOUNT_ID and CLOUDFLARE_API_TOKEN are not set; without them this stage does not exist.");
    process.exit(1);
  }
  const index = await buildIndex(LAWS, embed);
  const [vector] = await embed([question]);
  for (const hit of nearest(index, vector, NEAR_LAWS, 0)) {
    const law = LAWS.find((each) => each.id === hit.id);
    console.log(`${hit.score.toFixed(3)} ${hit.score >= NEAR_ENOUGH ? " " : "x"} ${hit.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);
  });
}