diff --git a/src/LLMProviders/chainRunner/utils/ThinkBlockStreamer.ts b/src/LLMProviders/chainRunner/utils/ThinkBlockStreamer.ts index 928104d4..a77c687a 100644 --- a/src/LLMProviders/chainRunner/utils/ThinkBlockStreamer.ts +++ b/src/LLMProviders/chainRunner/utils/ThinkBlockStreamer.ts @@ -9,6 +9,7 @@ import { createAIMessageWithToolCalls, } from "./nativeToolCalling"; import { logInfo, logWarn } from "@/logger"; +import { stripSpecialTokens } from "@/utils/stripSpecialTokens"; /** * ThinkBlockStreamer handles streaming content from various LLM providers @@ -116,7 +117,7 @@ export class ThinkBlockStreamer { private handleDeepseekChunk(chunk: any) { // Handle standard string content if (typeof chunk.content === "string") { - this.fullResponse += chunk.content; + this.fullResponse += stripSpecialTokens(chunk.content); } // Handle deepseek reasoning/thinking content diff --git a/src/utils/stripSpecialTokens.test.ts b/src/utils/stripSpecialTokens.test.ts new file mode 100644 index 00000000..b8d17d3b --- /dev/null +++ b/src/utils/stripSpecialTokens.test.ts @@ -0,0 +1,131 @@ +import { stripSpecialTokens } from "@/utils/stripSpecialTokens"; + +describe("stripSpecialTokens", () => { + // --- Individual token stripping --- + + it("strips ChatML <|im_end|>", () => { + expect(stripSpecialTokens("hello<|im_end|>")).toBe("hello"); + }); + + it("strips ChatML <|im_start|>", () => { + expect(stripSpecialTokens("<|im_start|>assistant")).toBe("assistant"); + }); + + it("strips Llama 3 <|eot_id|>", () => { + expect(stripSpecialTokens("done<|eot_id|>")).toBe("done"); + }); + + it("strips Llama 3 <|start_header_id|>", () => { + expect(stripSpecialTokens("<|start_header_id|>user")).toBe("user"); + }); + + it("strips Llama 3 <|end_header_id|>", () => { + expect(stripSpecialTokens("assistant<|end_header_id|>")).toBe("assistant"); + }); + + it("strips Gemma ", () => { + expect(stripSpecialTokens("response")).toBe("response"); + }); + + it("strips Gemma ", () => { + expect(stripSpecialTokens("model")).toBe("model"); + }); + + it("strips Phi <|end|>", () => { + expect(stripSpecialTokens("text<|end|>")).toBe("text"); + }); + + it("strips Phi <|assistant|>", () => { + expect(stripSpecialTokens("<|assistant|>answer")).toBe("answer"); + }); + + it("strips Phi <|user|>", () => { + expect(stripSpecialTokens("<|user|>question")).toBe("question"); + }); + + it("strips Phi <|system|>", () => { + expect(stripSpecialTokens("<|system|>prompt")).toBe("prompt"); + }); + + it("strips Mistral ", () => { + expect(stripSpecialTokens("end")).toBe("end"); + }); + + it("strips Mistral [INST]", () => { + expect(stripSpecialTokens("[INST]input")).toBe("input"); + }); + + it("strips Mistral [/INST]", () => { + expect(stripSpecialTokens("output[/INST]")).toBe("output"); + }); + + it("strips Qwen <|endoftext|>", () => { + expect(stripSpecialTokens("text<|endoftext|>")).toBe("text"); + }); + + it("strips DeepSeek <|end▁of▁sentence|>", () => { + expect(stripSpecialTokens("sentence<|end\u2581of\u2581sentence|>")).toBe("sentence"); + }); + + it("strips Command R <|END_OF_TURN_TOKEN|>", () => { + expect(stripSpecialTokens("turn<|END_OF_TURN_TOKEN|>")).toBe("turn"); + }); + + it("strips Command R <|START_OF_TURN_TOKEN|>", () => { + expect(stripSpecialTokens("<|START_OF_TURN_TOKEN|>next")).toBe("next"); + }); + + // --- Normal text is unchanged --- + + it("leaves normal text unchanged", () => { + const text = "This is a perfectly normal response with no special tokens."; + expect(stripSpecialTokens(text)).toBe(text); + }); + + it("leaves empty string unchanged", () => { + expect(stripSpecialTokens("")).toBe(""); + }); + + it("leaves text with regular angle brackets unchanged", () => { + const html = "
Hello world
"; + expect(stripSpecialTokens(html)).toBe(html); + }); + + it("does NOT strip (can appear in normal text)", () => { + // alone is not in the strip list; is (Mistral EOS token) and will be removed + expect(stripSpecialTokens("beginning")).toBe("beginning"); + }); + + it("strips Mistral EOS even when preceded by HTML-looking ", () => { + // is always stripped as it is the Mistral end-of-sequence token + expect(stripSpecialTokens("The strikethrough text here.")).toBe( + "The strikethrough text here." + ); + }); + + // --- Mixed content --- + + it("strips token at end of real text without affecting the rest", () => { + expect(stripSpecialTokens("Here is the answer.<|im_end|>")).toBe("Here is the answer."); + }); + + it("strips token at start of real text without affecting the rest", () => { + expect(stripSpecialTokens("<|im_start|>Here is the answer.")).toBe("Here is the answer."); + }); + + it("strips multiple tokens from a single chunk", () => { + expect(stripSpecialTokens("<|im_start|>assistant\nHello there!<|im_end|>")).toBe( + "assistant\nHello there!" + ); + }); + + it("strips tokens from different families in one pass", () => { + expect(stripSpecialTokens("[INST]question[/INST]<|im_end|>answer<|eot_id|>")).toBe( + "questionanswer" + ); + }); + + it("handles repeated occurrences of the same token", () => { + expect(stripSpecialTokens("<|im_end|>text<|im_end|>more<|im_end|>")).toBe("textmore"); + }); +}); diff --git a/src/utils/stripSpecialTokens.ts b/src/utils/stripSpecialTokens.ts new file mode 100644 index 00000000..8f3e96a9 --- /dev/null +++ b/src/utils/stripSpecialTokens.ts @@ -0,0 +1,68 @@ +/** + * Strips known chat template control tokens that local model servers (LM Studio, Ollama) + * sometimes leak into visible responses. These tokens are part of the model's internal + * chat template and should never appear in user-visible output. + * + * Token families covered: + * - ChatML: <|im_end|>, <|im_start|> + * - Llama 3: <|eot_id|>, <|start_header_id|>, <|end_header_id|> + * - Gemma: , + * - Phi: <|end|>, <|assistant|>, <|user|>, <|system|> + * - Mistral/Generic: , [INST], [/INST] + * - Qwen: <|endoftext|> + * - DeepSeek: <|end▁of▁sentence|> + * - Command R: <|END_OF_TURN_TOKEN|>, <|START_OF_TURN_TOKEN|> + */ + +/** Known special tokens that must never appear in user-visible output. */ +const SPECIAL_TOKEN_PATTERNS: string[] = [ + // ChatML + "<|im_end|>", + "<|im_start|>", + // Llama 3 + "<|eot_id|>", + "<|start_header_id|>", + "<|end_header_id|>", + // Gemma + "", + "", + // Phi + "<|end|>", + "<|assistant|>", + "<|user|>", + "<|system|>", + // Mistral/Generic (NOT — can appear in normal text) + "", + "[INST]", + "[/INST]", + // Qwen + "<|endoftext|>", + // DeepSeek (uses Unicode U+2581 LOWER ONE EIGHTH BLOCK for underscores) + "<|end\u2581of\u2581sentence|>", + // Command R + "<|END_OF_TURN_TOKEN|>", + "<|START_OF_TURN_TOKEN|>", +]; + +/** + * Pre-compiled regex built from all known special token patterns. + * Each token is escaped so that regex metacharacters are treated literally. + * The `g` flag ensures all occurrences in a chunk are removed in one pass. + */ +const SPECIAL_TOKENS_REGEX: RegExp = new RegExp( + SPECIAL_TOKEN_PATTERNS.map((t) => t.replace(/[.*+?^${}()|[\]\\]/g, "\\$&")).join("|"), + "g" +); + +/** + * Strips all known chat template control tokens from the provided text. + * + * This is a zero-cost operation when the input contains no special tokens + * (the regex engine scans without allocating a new string). + * + * @param text - Raw text chunk from a local LLM response stream. + * @returns The text with all known special tokens removed. + */ +export function stripSpecialTokens(text: string): string { + return text.replace(SPECIAL_TOKENS_REGEX, ""); +}