feat(extension): rank large Jev candidate lists with BM25

packages/extension/services/jevAct/bm25.ts

@@ -0,0 +1,52 @@
+import { describeCandidate, flattenText, tokens, type OutlineNode } from "./tree.js";
+import type { JsonValue } from "./typesafeClient.js";
+
+/** BM25 with the current ranker's name, text, and table weights retained. */
+export function scoreCandidatesBm25(
+  nodes: OutlineNode[],
+  candidates: OutlineNode[],
+  instruction: string,
+): Map<string, number> {
+  const query = new Set(tokens(instruction));
+  const documents = candidates.map((candidate) => {
+    const description = describeCandidate(nodes, candidate);
+    const fields: Array<[JsonValue | undefined, number]> = [
+      [description.name, 3],
+      [description.text, 2],
+      [description.within, 1],
+      [description.under_heading, 1],
+      [description.near_text, 1],
+      [description.group_text, 1],
+      [description.table, 2],
+    ];
+    const frequencies = new Map<string, number>();
+    let length = 0;
+    for (const [value, weight] of fields) {
+      for (const word of tokens(value === undefined ? "" : flattenText(value))) {
+        length += weight;
+        if (query.has(word)) frequencies.set(word, (frequencies.get(word) ?? 0) + weight);
+      }
+    }
+    return { id: candidate.id, frequencies, length };
+  });
+  const averageLength =
+    documents.reduce((sum, document) => sum + document.length, 0) / (documents.length || 1) || 1;
+  const documentFrequency = new Map<string, number>();
+  for (const document of documents) {
+    for (const word of document.frequencies.keys())
+      documentFrequency.set(word, (documentFrequency.get(word) ?? 0) + 1);
+  }
+  const scores = new Map<string, number>();
+  const n = documents.length;
+  for (const document of documents) {
+    let score = 0;
+    for (const [word, frequency] of document.frequencies) {
+      const df = documentFrequency.get(word)!;
+      const idf = Math.log(1 + (n - df + 0.5) / (df + 0.5));
+      const denominator = frequency + 1.2 * (0.65 + 0.35 * (document.length / averageLength));
+      score += idf * ((frequency * 2.2) / denominator);
+    }
+    scores.set(document.id, score);
+  }
+  return scores;
+}

packages/extension/services/jevAct/pick.ts

@@ -1,11 +1,11 @@
 import type { StagehandLogger } from "../../logger.js";
 import { quotedStrings } from "./args.js";
+import { scoreCandidatesBm25 } from "./bm25.js";
 import {
   buildView,
   describeCandidate,
   exactNameMatches,
   nearbyTwins,
-  scoreCandidates,
   type OutlineNode,
   type ViewKind,
 } from "./tree.js";
@@ -183,7 +183,7 @@ export async function pickTarget(
 
     const attempts: OutlineNode[][] = [];
     if (candidates.length > PRUNE_ABOVE) {
-      const scores = scoreCandidates(snap.nodes, candidates, ctx.instruction);
+      const scores = scoreCandidatesBm25(snap.nodes, candidates, ctx.instruction);
       const pruned = candidates
         .filter((candidate) => (scores.get(candidate.id) ?? 0) > 0)
         .sort((a, b) => (scores.get(b.id) ?? 0) - (scores.get(a.id) ?? 0) || a.index - b.index)

packages/extension/services/jevAct/tree.ts

@@ -603,14 +603,14 @@ const STOPWORDS = new Set(
 );
 
 /** Values only: object keys ("row", "column") are ours, not the page's. */
-function flattenText(value: JsonValue): string {
+export function flattenText(value: JsonValue): string {
   if (typeof value === "string") return value;
   if (Array.isArray(value)) return value.map(flattenText).join(" ");
   if (value && typeof value === "object") return Object.values(value).map(flattenText).join(" ");
   return value === null ? "" : String(value);
 }
 
-function tokens(text: string): string[] {
+export function tokens(text: string): string[] {
   return (
     (text.toLowerCase().match(/[\p{L}\p{N}]+/gu) ?? [])
       // "15th" should find the calendar cell named "15".

packages/extension/tests/jevActModules.test.ts

@@ -8,6 +8,7 @@ import {
   matchOption,
 } from "../services/jevAct/args.js";
 import { blockingSignal, readPageState } from "../services/jevAct/pageState.js";
+import { scoreCandidatesBm25 } from "../services/jevAct/bm25.js";
 import type { AskContext } from "../services/jevAct/pick.js";
 import {
   describeCandidate,
@@ -77,6 +78,35 @@ describe("jev act candidate helpers", () => {
     expect(scoreCandidates(nodes, buttons, "open the basket").get("0-6")).toBe(0);
   });
 
+  it("keeps a rare row name in the top 30 despite repeated generic terms", () => {
+    const lines = ["[0-1] RootWebArea: Records", "  [0-2] table: Records"];
+    let id = 3;
+    for (const name of [
+      ...Array.from({ length: 60 }, (_, index) => `Account ${index + 1}`),
+      "Acme",
+      ...Array.from({ length: 20 }, (_, index) => `Other ${index + 1}`),
+    ]) {
+      lines.push(`    [0-${id++}] row`);
+      lines.push(`      [0-${id++}] cell: ${name}`);
+      lines.push(`      [0-${id++}] button: Delete`);
+    }
+    const nodes = parseOutline(lines.join("\n"));
+    const buttons = nodes.filter((node) => node.role === "button");
+    const ranked = (scores: Map<string, number>) =>
+      buttons
+        .filter((button) => (scores.get(button.id) ?? 0) > 0)
+        .sort((a, b) => (scores.get(b.id) ?? 0) - (scores.get(a.id) ?? 0) || a.index - b.index)
+        .slice(0, 30)
+        .map((button) => button.id);
+    expect(ranked(scoreCandidates(nodes, buttons, "Delete the Acme account"))).not.toContain(
+      "0-185",
+    );
+    expect(ranked(scoreCandidatesBm25(nodes, buttons, "Delete the Acme account"))[0]).toBe("0-185");
+    expect([...scoreCandidatesBm25(nodes, buttons, "Erase the obsolete entry").values()]).toEqual(
+      Array(buttons.length).fill(0),
+    );
+  });
+
   it("matches quoted names exactly and case-insensitively", () => {
     const nodes = parseOutline(PAGE);
     expect(exactNameMatches(nodes, "cart").map((node) => node.id)).toEqual(["0-3"]);

scripts/compare-jev-rankers.ts

@@ -0,0 +1,137 @@
+import { performance } from "node:perf_hooks";
+import { scoreCandidatesBm25 } from "../packages/extension/services/jevAct/bm25.js";
+import {
+  buildView,
+  parseOutline,
+  scoreCandidates,
+} from "../packages/extension/services/jevAct/tree.js";
+
+// Run with `node_modules/.bin/tsx scripts/compare-jev-rankers.ts [--jev]`.
+// Generated outlines isolate ranking behavior; they are not a sample of live pages.
+// --jev asks TypeSafe about both shortlists in the rare-name case.
+
+type Case = { name: string; instruction: string; outline: string; target: string };
+
+function rows(names: string[], button = "Delete"): string {
+  const lines = ["[0-1] RootWebArea: Records", "  [0-2] table: Records"];
+  let id = 3;
+  for (const name of names) {
+    lines.push(`    [0-${id++}] row`);
+    lines.push(`      [0-${id++}] cell: ${name}`);
+    lines.push(`      [0-${id++}] button: ${button}`);
+  }
+  return lines.join("\n");
+}
+
+const cases: Case[] = [
+  {
+    name: "invoice row among 80 rows",
+    instruction: "Delete invoice 1042",
+    outline: rows(Array.from({ length: 80 }, (_, i) => `Invoice ${1000 + i}`)),
+    target: "0-131",
+  },
+  {
+    name: "rare account name after 60 common-name rows",
+    instruction: "Delete the Acme account",
+    outline: rows([
+      ...Array.from({ length: 60 }, (_, i) => `Account ${i + 1}`),
+      "Acme",
+      ...Array.from({ length: 20 }, (_, i) => `Other ${i + 1}`),
+    ]),
+    target: "0-185",
+  },
+  {
+    name: "long row with target name",
+    instruction: "Delete the red toaster",
+    outline: rows([
+      ...Array.from({ length: 55 }, (_, i) => `Red product ${i + 1}`),
+      "Red toaster with an unusually long product description and details",
+      ...Array.from({ length: 24 }, (_, i) => `Blue product ${i + 1}`),
+    ]),
+    target: "0-170",
+  },
+  {
+    name: "semantic mismatch",
+    instruction: "Erase the obsolete entry",
+    outline: rows(Array.from({ length: 80 }, (_, i) => `Archive ${i + 1}`)),
+    target: "0-5",
+  },
+];
+
+const live = process.argv.includes("--jev");
+if (live) {
+  const dotenv = await import("dotenv");
+  dotenv.config({ quiet: true });
+}
+
+function shortlist(scores: Map<string, number>, candidates: ReturnType<typeof buildView>) {
+  return candidates
+    .filter((candidate) => (scores.get(candidate.id) ?? 0) > 0)
+    .sort((a, b) => (scores.get(b.id) ?? 0) - (scores.get(a.id) ?? 0) || a.index - b.index)
+    .slice(0, 30);
+}
+
+for (const testCase of cases) {
+  const nodes = parseOutline(testCase.outline);
+  const candidates = buildView(nodes, "pointer");
+  const baseline = shortlist(scoreCandidates(nodes, candidates, testCase.instruction), candidates);
+  const bm25 = shortlist(scoreCandidatesBm25(nodes, candidates, testCase.instruction), candidates);
+  const position = (list: typeof candidates) =>
+    list.findIndex((item) => item.id === testCase.target) + 1;
+  const time = (score: typeof scoreCandidates) => {
+    const start = performance.now();
+    for (let i = 0; i < 100; i++) score(nodes, candidates, testCase.instruction);
+    return (performance.now() - start) / 100;
+  };
+  process.stdout.write(
+    `${JSON.stringify({
+      case: testCase.name,
+      candidates: candidates.length,
+      baselineRank: position(baseline),
+      bm25Rank: position(bm25),
+      baselineMs: Number(time(scoreCandidates).toFixed(3)),
+      bm25Ms: Number(time(scoreCandidatesBm25).toFixed(3)),
+    })}\n`,
+  );
+  if (live && testCase.name === "rare account name after 60 common-name rows") {
+    const { describeCandidate } = await import("../packages/extension/services/jevAct/tree.js");
+    const { choiceAnswer, systemOne } =
+      await import("../packages/extension/services/jevAct/typesafeClient.js");
+    const apiKey = process.env.TYPESAFE_API_KEY;
+    if (!apiKey) throw new Error("TYPESAFE_API_KEY is missing");
+    for (const [ranker, list] of [
+      ["baseline", baseline],
+      ["bm25", bm25],
+    ] as const) {
+      const criteria = Object.fromEntries(
+        list.map((candidate) => [candidate.id, describeCandidate(nodes, candidate)]),
+      );
+      const response = await systemOne(
+        { apiKey },
+        { instruction: testCase.instruction },
+        {
+          best: {
+            type: "choice",
+            instructions: "Which element does the instruction refer to?",
+            criteria,
+          },
+          strict: {
+            type: "choice",
+            instructions: "Which element does the instruction refer to?",
+            criteria: { ...criteria, none_match: "None of these elements matches the instruction" },
+          },
+        },
+      );
+      process.stdout.write(
+        `${JSON.stringify({
+          ranker,
+          jevChoice: choiceAnswer(response, "best").choice,
+          jevConfidence: choiceAnswer(response, "best").confidence,
+          jevNone: choiceAnswer(response, "strict").probabilities.none_match,
+          jevMs: response.durationMs,
+          inputTokens: response.usage.inputTokens,
+        })}\n`,
+      );
+    }
+  }
+}