From d3083c7117ecc72f73b0f1cb6abb09c370a3b574 Mon Sep 17 00:00:00 2001 From: lorenzozanee Date: Wed, 16 Sep 2026 01:11:59 +0800 Subject: [PATCH] fix(compression): preserve case, tags and negations through llmlingua engine --- .../compression/engines/llmlingua/index.ts | 84 ++++++++- .../compression/llmlingua-fidelity.test.ts | 161 ++++++++++++++++++ 2 files changed, 240 insertions(+), 5 deletions(-) create mode 100644 tests/unit/compression/llmlingua-fidelity.test.ts diff --git a/open-sse/services/compression/engines/llmlingua/index.ts b/open-sse/services/compression/engines/llmlingua/index.ts index 83e12e716800..1dea61227703 100644 --- a/open-sse/services/compression/engines/llmlingua/index.ts +++ b/open-sse/services/compression/engines/llmlingua/index.ts @@ -20,7 +20,10 @@ * preserved constructs) into placeholder strings. The prose between placeholders * is what gets sent to the backend. Code blocks are re-stitched verbatim into * the output. This is done REGARDLESS of what the backend does — the engine - * physically never passes code to the model. + * physically never passes code to the model. XML tags and negation/absolute + * words are additionally split out of prose (the uncased default model breaks + * tag punctuation and prunes negations), and backend replies are mapped back + * to the original casing/edge whitespace before re-stitching. * * ### Fail-open points (all errors → original body) * 1. Backend rejects for a prose segment → catch → segment kept as-is. @@ -94,6 +97,61 @@ interface TextSegment { text: string; } +/** + * Spans the backend must never see, split out of prose before the call and + * re-stitched verbatim afterwards: + * - XML-style tags (``, ``, `
`, ...). The uncased default + * model rebuilds output from word pieces, which breaks tag punctuation; + * lone tags outside the known preserved envelopes need the same shield. + * - Negations/absolutes (`not`, `never`, `must`, `don't`, ...). Pruning one + * inverts an instruction, so they are not prunable by construction. + * Hyphenated compounds (`no-op`) are excluded via the lookarounds so ordinary + * words are never fragmented. + */ +const PROTECTED_SPAN_RE = + /(<\/?[A-Za-z][A-Za-z0-9._-]*(?:\s[^<>]*?)?\/?>|\b\w+(?:n't|n’t)\b|(? cursor) out.push({ kind: "prose", text: prose.slice(cursor, idx) }); + out.push({ kind: "preserved", text: m[0] }); + cursor = idx + m[0].length; + } + if (cursor < prose.length) out.push({ kind: "prose", text: prose.slice(cursor) }); + return out; +} + +/** Word token for case restoration (keeps `don't`-style contractions whole). */ +const WORD_RE = /[A-Za-z0-9]+(?:['\u2019][A-Za-z0-9]+)*/g; + +/** + * Map surviving backend words back to their original character spans. + * The default model is uncased, so its output is lowercased prose rebuilt from + * word pieces; for each output word, restore the casing of the matching source + * word (first unused case-insensitive match wins, in order). Words the backend + * invented match nothing and are left as-is. Whole-word only. + */ +function restoreSourceCase(source: string, output: string): string { + const forms = new Map(); + for (const w of source.matchAll(WORD_RE)) { + const key = w[0].toLowerCase(); + const list = forms.get(key); + if (list) list.push(w[0]); + else forms.set(key, [w[0]]); + } + if (forms.size === 0) return output; + return output.replace(WORD_RE, (w) => { + const list = forms.get(w.toLowerCase()); + const form = list?.shift(); + return form ?? w; + }); +} + /** * Split `text` into alternating prose / preserved segments using * `extractPreservedBlocks` from preservation.ts. @@ -111,7 +169,7 @@ function splitProseAndPreserved(text: string): TextSegment[] { const { text: withPlaceholders, blocks } = extractPreservedBlocks(text); if (blocks.length === 0) { - return [{ kind: "prose", text }]; + return splitProtectedSpans(text); } const segments: TextSegment[] = []; @@ -129,7 +187,8 @@ function splitProseAndPreserved(text: string): TextSegment[] { if (original !== undefined) { segments.push({ kind: "preserved", text: original }); } else { - segments.push({ kind: "prose", text: part }); + // Shield tags + negations/absolutes from the backend before it sees them. + segments.push(...splitProtectedSpans(part)); } } @@ -147,6 +206,11 @@ type MessageLike = { /** * Compress a single prose string via the backend. * On any error, fail-open and return the original text. + * + * The uncased default model lowercases and re-spaces its output, so a raw + * backend reply is never stitched in directly: surviving words are mapped back + * to their original spans (casing) and edge whitespace the backend ate is + * restored, keeping newlines/spacing around preserved segments intact. */ async function compressProseText( text: string, @@ -156,9 +220,19 @@ async function compressProseText( if (!text.trim()) return { text, didCompress: false }; try { const compressed = await backend(text, opts); + // An empty reply carries no surviving words — stitching it in would delete + // the segment, so treat it like any other backend failure. + if (typeof compressed !== "string" || !compressed.trim()) { + return { text, didCompress: false }; + } + let out = restoreSourceCase(text, compressed); + const leading = text.match(/^\s*/)?.[0] ?? ""; + const trailing = text.match(/\s*$/)?.[0] ?? ""; + if (leading && !/^\s/.test(out)) out = leading + out; + if (trailing && !/\s$/.test(out)) out = out + trailing; // Accept only if it actually gets shorter (reject no-ops or expansions) - if (typeof compressed === "string" && compressed.length < text.length) { - return { text: compressed, didCompress: true }; + if (out.length < text.length) { + return { text: out, didCompress: true }; } return { text, didCompress: false }; } catch { diff --git a/tests/unit/compression/llmlingua-fidelity.test.ts b/tests/unit/compression/llmlingua-fidelity.test.ts new file mode 100644 index 000000000000..2fb82f022b62 --- /dev/null +++ b/tests/unit/compression/llmlingua-fidelity.test.ts @@ -0,0 +1,161 @@ +/** + * Fidelity tests for the llmlingua engine (#13455). + * + * The default TinyBERT model is uncased (`do_lower_case: true`): it rebuilds + * output from lower-cased word pieces, prunes negations/absolutes and eats the + * whitespace around preserved spans. The engine must therefore shield tags + + * negations from the backend, map surviving words back to their original spans + * (casing) and restore edge whitespace on re-stitch — all model-free, via the + * `setLlmlinguaBackend` hook. + */ + +import { describe, it, after } from "node:test"; +import assert from "node:assert/strict"; + +import { + llmlinguaEngine, + setLlmlinguaBackend, +} from "../../../open-sse/services/compression/engines/llmlingua/index.ts"; + +function makeBody(content: string): Record { + return { model: "gpt-4o", messages: [{ role: "user", content }] }; +} + +function outContent(result: Awaited>): string { + return (result.body.messages as Array<{ role: string; content: string }>)[0]!.content; +} + +/** + * Simulates the uncased-model backend: lowercases, trims edges, collapses + * inner whitespace runs and drops a negation — i.e. what the real TinyBERT + * path observably does to prose. + */ +function uncasedModelBackend(text: string): Promise { + let out = text.toLowerCase(); + out = out.replace(/\bdon't\b/g, ""); + out = out.trim().replace(/\s+/g, " "); + return Promise.resolve(out); +} + +after(() => { + setLlmlinguaBackend(null); +}); + +describe("llmlingua engine — fidelity (#13455)", () => { + it("case, newlines and spacing around preserved spans survive re-stitching", async () => { + setLlmlinguaBackend(uncasedModelBackend); + + const code = "`rm -rf`"; + const content = `Deploy rules:\n\nNever run this:\n\n${code}\n\nAlways verify first.`; + const result = await llmlinguaEngine.applyAsync!(makeBody(content), { + stepConfig: { minTokens: 0 }, + }); + + const out = outContent(result); + // Preserved span byte-identical, with original newlines/spacing around it. + assert.ok( + out.includes(`this:\n\n${code}\n\nAlways`), + `newlines/spacing around ${code} must survive.\nGot:\n${out}` + ); + // Surviving words keep their original casing despite the lower-casing backend. + assert.ok(out.includes("Deploy rules:"), `prose casing must be restored.\nGot:\n${out}`); + // The absolutes Never/Always are shielded verbatim (original case). + assert.ok(out.includes("Never run this:"), `negation casing must survive.\nGot:\n${out}`); + }); + + it("negations/absolutes never reach the backend and are kept verbatim", async () => { + const seen: string[] = []; + setLlmlinguaBackend((text) => { + seen.push(text); + return uncasedModelBackend(text); + }); + + const content = + "NEVER run rm on the target host. Do not push to main directly. " + + "You must ask first. Don't restart the database."; + const result = await llmlinguaEngine.applyAsync!(makeBody(content), { + stepConfig: { minTokens: 0 }, + }); + + const joined = seen.join("\n").toLowerCase(); + for (const word of ["never", "not", "must", "don't"]) { + assert.ok( + !new RegExp(`\\b${word}\\b`).test(joined), + `backend must never see ${word}; got backend inputs:\n${seen.join("\n")}` + ); + } + + const out = outContent(result); + for (const word of ["NEVER", "Do not", "must", "Don't"]) { + assert.ok(out.includes(word), `output must keep ${word} verbatim.\nGot:\n${out}`); + } + }); + + it("XML-style tags are never sent to the backend and survive verbatim", async () => { + const seen: string[] = []; + setLlmlinguaBackend((text) => { + seen.push(text); + // Backend that additionally strips anything looking like a tag. + return uncasedModelBackend(text).then((s) => s.replace(/<[^>]*>/g, "")); + }); + + const content = "Hello world, see bar and a lone here."; + const result = await llmlinguaEngine.applyAsync!(makeBody(content), { + stepConfig: { minTokens: 0 }, + }); + + for (const call of seen) { + assert.ok(!call.includes("<"), `backend must not receive tags, got:\n${call}`); + } + + const out = outContent(result); + for (const tag of ["", "", "", ""]) { + assert.ok(out.includes(tag), `output must keep ${tag} verbatim.\nGot:\n${out}`); + } + }); + + it("system-reminder envelope stays byte-identical and keeps boundary newlines", async () => { + setLlmlinguaBackend(uncasedModelBackend); + + const envelope = + "\nProject instructions:\n\n- NEVER run `rm -rf`.\n- Do NOT edit `/etc`.\n"; + const content = `${envelope}\n\nPlease deploy the fix. Thanks!`; + const result = await llmlinguaEngine.applyAsync!(makeBody(content), { + stepConfig: { minTokens: 0 }, + }); + + const out = outContent(result); + assert.ok(out.includes(envelope), `envelope must be byte-identical.\nGot:\n${out}`); + assert.ok( + out.includes("\n\nPlease"), + `blank line after the envelope must survive (was: please).\nGot:\n${out}` + ); + }); + + it("empty backend output fail-opens instead of deleting the segment", async () => { + setLlmlinguaBackend(() => Promise.resolve("")); + + const content = "Some prose the backend wipes out entirely."; + const body = makeBody(content); + const result = await llmlinguaEngine.applyAsync!(body, { stepConfig: { minTokens: 0 } }); + + assert.equal(result.compressed, false, "empty reply must fail-open"); + assert.equal(outContent(result), content, "original segment must be kept"); + }); + + it("hyphenated compounds like no-op are not fragmented", async () => { + const seen: string[] = []; + setLlmlinguaBackend((text) => { + seen.push(text); + return Promise.resolve("X"); + }); + + const content = "This change is a no-op for existing users."; + await llmlinguaEngine.applyAsync!(makeBody(content), { stepConfig: { minTokens: 0 } }); + + assert.ok( + seen.some((call) => call.includes("no-op")), + `no-op must reach the backend whole; got backend inputs:\n${seen.join("\n")}` + ); + }); +});