ask/fuse.ts

import type { Law } from "./corpus.js";
import { RETRIEVED_LAWS, bm25 } from "./retrieve.js";

// Two rankings over the same laws, merged. Neither ranking's raw scores are
// compared with the other's, because they are not on the same scale; what is
// compared is where each law stands in each list.

/** The usual constant for reciprocal-rank fusion: it keeps one list's first place from swamping the other's. */
export const RRF_K = 60;

/** The question again in the law's words (Chapter 6), and what the index found for it. */
export interface Search {
  search: string;
  near: string[];
}

/** One list from several, by reciprocal rank: a law near the top of any list makes the shortlist, and one high on all of them leads it. */
export function fuse(rankings: string[][], k = RRF_K): string[] {
  const fused = new Map<string, number>();
  for (const ids of rankings) {
    ids.forEach((id, rank) => fused.set(id, (fused.get(id) ?? 0) + 1 / (k + rank + 1)));
  }
  return [...fused.entries()].sort((a, b) => b[1] - a[1]).map(([id]) => id);
}

/**
 * The laws most likely to answer `question`, best first: BM25 over the
 * question, the ids the index found nearest it, and each rephrased search's
 * own two rankings, fused. With no `near` and no `searches`, this is BM25
 * alone, byte for byte.
 */
export function retrieve(
  laws: Law[],
  question: string,
  limit = RETRIEVED_LAWS,
  near: string[] = [],
  searches: Search[] = [],
): Law[] {
  const byId = new Map(laws.map((law) => [law.id, law]));
  // Only ids the file holds: a vector left behind for a law since removed is
  // dropped here, never shown.
  const meant = near.filter((id) => byId.has(id));
  if (meant.length === 0 && searches.length === 0) return bm25(laws, question, limit);
  const rankings = [bm25(laws, question, limit * 2).map((law) => law.id), meant];
  for (const { search, near: nearer } of searches) {
    rankings.push(bm25(laws, search, limit * 2).map((law) => law.id));
    rankings.push(nearer.filter((id) => byId.has(id)));
  }
  return fuse(rankings)
    .slice(0, limit)
    .map((id) => byId.get(id) as Law);
}