From 0ad3f8632f9f0db9d007c12cb3b4bd0cae09e8e4 Mon Sep 17 00:00:00 2001 From: Vishal Anton Date: Thu, 17 Sep 2026 16:11:34 +0530 Subject: [PATCH] feat(extension): rank large Jev candidate lists with BM25 --- packages/extension/services/jevAct/bm25.ts | 52 +++++++ packages/extension/services/jevAct/pick.ts | 4 +- packages/extension/services/jevAct/tree.ts | 4 +- .../extension/tests/jevActModules.test.ts | 30 ++++ scripts/compare-jev-rankers.ts | 137 ++++++++++++++++++ 5 files changed, 223 insertions(+), 4 deletions(-) create mode 100644 packages/extension/services/jevAct/bm25.ts create mode 100644 scripts/compare-jev-rankers.ts diff --git a/packages/extension/services/jevAct/bm25.ts b/packages/extension/services/jevAct/bm25.ts new file mode 100644 index 0000000000..3b81a70f71 --- /dev/null +++ b/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 { + 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(); + 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(); + for (const document of documents) { + for (const word of document.frequencies.keys()) + documentFrequency.set(word, (documentFrequency.get(word) ?? 0) + 1); + } + const scores = new Map(); + 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; +} diff --git a/packages/extension/services/jevAct/pick.ts b/packages/extension/services/jevAct/pick.ts index 72a1acca7c..f1a21166a9 100644 --- a/packages/extension/services/jevAct/pick.ts +++ b/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) diff --git a/packages/extension/services/jevAct/tree.ts b/packages/extension/services/jevAct/tree.ts index 7a84fcc8c9..cc2b25033c 100644 --- a/packages/extension/services/jevAct/tree.ts +++ b/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". diff --git a/packages/extension/tests/jevActModules.test.ts b/packages/extension/tests/jevActModules.test.ts index 894d75c3d5..6c2484d7d2 100644 --- a/packages/extension/tests/jevActModules.test.ts +++ b/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) => + 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"]); diff --git a/scripts/compare-jev-rankers.ts b/scripts/compare-jev-rankers.ts new file mode 100644 index 0000000000..b9a61eda18 --- /dev/null +++ b/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, candidates: ReturnType) { + 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`, + ); + } + } +}