fix(compression): preserve case, tags and negations through llmlingua engine

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 (`<tag>`, `</tag>`, `<br/>`, ...). 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|(?<![A-Za-z0-9_-])(?:not|no|never|none|nothing|nobody|nowhere|neither|nor|cannot|always|must)(?![A-Za-z0-9_-]))/gi;
+
+/** Split protected spans out of a prose segment; they re-stitch verbatim. */
+function splitProtectedSpans(prose: string): TextSegment[] {
+  PROTECTED_SPAN_RE.lastIndex = 0;
+  const out: TextSegment[] = [];
+  let cursor = 0;
+  for (const m of prose.matchAll(PROTECTED_SPAN_RE)) {
+    const idx = m.index ?? 0;
+    if (idx > 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<string, string[]>();
+  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 {

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<string, unknown> {
+  return { model: "gpt-4o", messages: [{ role: "user", content }] };
+}
+
+function outContent(result: Awaited<ReturnType<typeof llmlinguaEngine.applyAsync>>): 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<string> {
+  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 <tag> world, see <foo>bar</foo> and a lone </tag> 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 ["<tag>", "<foo>", "</foo>", "</tag>"]) {
+      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 =
+      "<system-reminder>\nProject instructions:\n\n- NEVER run `rm -rf`.\n- Do NOT edit `/etc`.\n</system-reminder>";
+    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("</system-reminder>\n\nPlease"),
+      `blank line after the envelope must survive (was: </system-reminder>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")}`
+    );
+  });
+});