mirror of
https://github.com/logancyang/obsidian-copilot.git
synced 2026-07-22 07:50:24 +00:00
Local model servers (LM Studio, Ollama) sometimes leak chat template control tokens like <|im_end|> into visible responses when stop tokens are not configured on the server side. Add stripSpecialTokens utility that removes known tokens (ChatML, Llama 3, Gemma, Phi, Mistral, Qwen, DeepSeek, Command R) via a single pre-compiled regex. Applied in ThinkBlockStreamer.handleDeepseekChunk before content is appended to the response. Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
parent
0ac60eb31e
commit
6d14b3ee89
3 changed files with 201 additions and 1 deletions
|
|
@ -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
|
||||
|
|
|
|||
131
src/utils/stripSpecialTokens.test.ts
Normal file
131
src/utils/stripSpecialTokens.test.ts
Normal file
|
|
@ -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 <end_of_turn>", () => {
|
||||
expect(stripSpecialTokens("response<end_of_turn>")).toBe("response");
|
||||
});
|
||||
|
||||
it("strips Gemma <start_of_turn>", () => {
|
||||
expect(stripSpecialTokens("<start_of_turn>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 </s>", () => {
|
||||
expect(stripSpecialTokens("end</s>")).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 = "<div>Hello <strong>world</strong></div>";
|
||||
expect(stripSpecialTokens(html)).toBe(html);
|
||||
});
|
||||
|
||||
it("does NOT strip <s> (can appear in normal text)", () => {
|
||||
// <s> alone is not in the strip list; </s> is (Mistral EOS token) and will be removed
|
||||
expect(stripSpecialTokens("<s>beginning")).toBe("<s>beginning");
|
||||
});
|
||||
|
||||
it("strips </s> Mistral EOS even when preceded by HTML-looking <s>", () => {
|
||||
// </s> is always stripped as it is the Mistral end-of-sequence token
|
||||
expect(stripSpecialTokens("The <s>strikethrough</s> text here.")).toBe(
|
||||
"The <s>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");
|
||||
});
|
||||
});
|
||||
68
src/utils/stripSpecialTokens.ts
Normal file
68
src/utils/stripSpecialTokens.ts
Normal file
|
|
@ -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: <end_of_turn>, <start_of_turn>
|
||||
* - Phi: <|end|>, <|assistant|>, <|user|>, <|system|>
|
||||
* - Mistral/Generic: </s>, [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
|
||||
"<end_of_turn>",
|
||||
"<start_of_turn>",
|
||||
// Phi
|
||||
"<|end|>",
|
||||
"<|assistant|>",
|
||||
"<|user|>",
|
||||
"<|system|>",
|
||||
// Mistral/Generic (NOT <s> — can appear in normal text)
|
||||
"</s>",
|
||||
"[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, "");
|
||||
}
|
||||
Loading…
Reference in a new issue