101 lines
3.7 KiB
TypeScript
101 lines
3.7 KiB
TypeScript
// ThinkSplitter — peels inline <think>...</think> reasoning out of streamed text
|
|
// content. Some local models (QwQ, DeepSeek-R1 distills, MiniMax) served raw
|
|
// emit their chain-of-thought inline in the assistant `content` channel rather
|
|
// than on a structured reasoning channel; BooCode's stream adapter otherwise
|
|
// treats that as ordinary prose. This splitter routes the reasoning span to the
|
|
// reasoning accumulator and passes the rest through unchanged.
|
|
//
|
|
// Ported from deepseek-reasonix internal/provider/openai/think.go. Two
|
|
// guarantees make it safe to run on every text delta:
|
|
// 1. It only ARMS if the turn's content begins with <think> (after leading
|
|
// whitespace), so an answer that merely mentions the tag is never hijacked.
|
|
// 2. For any content that does not start with <think> it degrades to a
|
|
// verbatim pass-through (a no-op for models on a structured reasoning
|
|
// channel).
|
|
// It buffers partial closing tags across chunk boundaries so a `</thi` split
|
|
// across two deltas is not mistaken for prose.
|
|
|
|
const OPEN = '<think>';
|
|
const CLOSE = '</think>';
|
|
const LEADING_WS = /^[ \t\r\n]+/;
|
|
|
|
type State = 'probe' | 'inside' | 'passthrough';
|
|
|
|
export interface SplitResult {
|
|
/** Text classified as reasoning (the inside of a <think> block). */
|
|
reasoning: string;
|
|
/** Text classified as ordinary content to pass through. */
|
|
text: string;
|
|
}
|
|
|
|
/**
|
|
* Longest proper suffix of `s` that is a prefix of `marker`. Used to hold back
|
|
* the bytes that might be the start of a closing tag split across chunks. Never
|
|
* returns the full marker length (that is a complete match, handled separately).
|
|
*/
|
|
function markerSuffixLen(s: string, marker: string): number {
|
|
const max = Math.min(marker.length - 1, s.length);
|
|
for (let n = max; n > 0; n--) {
|
|
if (marker.startsWith(s.slice(s.length - n))) return n;
|
|
}
|
|
return 0;
|
|
}
|
|
|
|
/** Stateful, single-stream splitter. Create one per streamed completion. */
|
|
export class ThinkSplitter {
|
|
private state: State = 'probe';
|
|
private buf = '';
|
|
|
|
push(s: string): SplitResult {
|
|
if (this.state === 'passthrough') return { reasoning: '', text: s };
|
|
if (this.state === 'inside') return this.scanClose(s);
|
|
|
|
// probe
|
|
this.buf += s;
|
|
const trimmed = this.buf.replace(LEADING_WS, '');
|
|
if (trimmed.length < OPEN.length) {
|
|
// Not enough yet to decide. Hold only if still a viable <think> prefix.
|
|
if (OPEN.startsWith(trimmed)) return { reasoning: '', text: '' };
|
|
return this.drainPassthrough();
|
|
}
|
|
if (trimmed.startsWith(OPEN)) {
|
|
this.state = 'inside';
|
|
this.buf = '';
|
|
return this.scanClose(trimmed.slice(OPEN.length));
|
|
}
|
|
return this.drainPassthrough();
|
|
}
|
|
|
|
/** Resolve any buffered remainder at stream end. */
|
|
flush(): SplitResult {
|
|
const r = this.buf;
|
|
this.buf = '';
|
|
if (this.state === 'inside') return { reasoning: r, text: '' };
|
|
return { reasoning: '', text: r };
|
|
}
|
|
|
|
private scanClose(s: string): SplitResult {
|
|
this.buf += s;
|
|
const idx = this.buf.indexOf(CLOSE);
|
|
if (idx >= 0) {
|
|
const reasoning = this.buf.slice(0, idx);
|
|
const text = this.buf.slice(idx + CLOSE.length).replace(LEADING_WS, '');
|
|
this.buf = '';
|
|
this.state = 'passthrough';
|
|
return { reasoning, text };
|
|
}
|
|
// No full closing tag yet — emit everything except a possible partial tag.
|
|
const keep = markerSuffixLen(this.buf, CLOSE);
|
|
const reasoning = this.buf.slice(0, this.buf.length - keep);
|
|
this.buf = this.buf.slice(this.buf.length - keep);
|
|
return { reasoning, text: '' };
|
|
}
|
|
|
|
private drainPassthrough(): SplitResult {
|
|
const text = this.buf;
|
|
this.buf = '';
|
|
this.state = 'passthrough';
|
|
return { reasoning: '', text };
|
|
}
|
|
}
|