Spaces:
Running
Running
File size: 2,908 Bytes
f672a5d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 | export interface ThinkDelta {
type: "reasoning" | "content";
textDelta: string;
}
export class ThinkStreamParser {
reasoning = "";
content = "";
private inThink = false;
private buffer = "";
private static readonly OPEN_TAG = "<think>";
private static readonly CLOSE_TAG = "</think>";
push(text: string): ThinkDelta[] {
const deltas: ThinkDelta[] = [];
this.buffer += text;
while (this.buffer.length > 0) {
if (this.inThink) {
const closeIndex = this.buffer.indexOf(ThinkStreamParser.CLOSE_TAG);
if (closeIndex !== -1) {
const chunk = this.buffer.slice(0, closeIndex);
if (chunk) {
this.reasoning += chunk;
deltas.push({ type: "reasoning", textDelta: chunk });
}
this.buffer = this.buffer.slice(
closeIndex + ThinkStreamParser.CLOSE_TAG.length,
);
this.inThink = false;
continue;
}
const safeLength = this.getSafeFlushLength(
this.buffer,
ThinkStreamParser.CLOSE_TAG,
);
if (safeLength > 0) {
const chunk = this.buffer.slice(0, safeLength);
this.reasoning += chunk;
deltas.push({ type: "reasoning", textDelta: chunk });
this.buffer = this.buffer.slice(safeLength);
}
break;
}
const openIndex = this.buffer.indexOf(ThinkStreamParser.OPEN_TAG);
if (openIndex !== -1) {
const chunk = this.buffer.slice(0, openIndex);
if (chunk) {
this.content += chunk;
deltas.push({ type: "content", textDelta: chunk });
}
this.buffer = this.buffer.slice(
openIndex + ThinkStreamParser.OPEN_TAG.length,
);
this.inThink = true;
continue;
}
const safeLength = this.getSafeFlushLength(
this.buffer,
ThinkStreamParser.OPEN_TAG,
);
if (safeLength > 0) {
const chunk = this.buffer.slice(0, safeLength);
this.content += chunk;
deltas.push({ type: "content", textDelta: chunk });
this.buffer = this.buffer.slice(safeLength);
}
break;
}
return deltas;
}
flush(): ThinkDelta[] {
if (!this.buffer) return [];
const deltas: ThinkDelta[] = [];
if (this.inThink) {
this.reasoning += this.buffer;
deltas.push({ type: "reasoning", textDelta: this.buffer });
} else {
this.content += this.buffer;
deltas.push({ type: "content", textDelta: this.buffer });
}
this.buffer = "";
return deltas;
}
private getSafeFlushLength(buffer: string, tag: string): number {
for (
let overlap = Math.min(buffer.length, tag.length - 1);
overlap > 0;
overlap--
) {
if (buffer.endsWith(tag.slice(0, overlap))) {
return buffer.length - overlap;
}
}
return buffer.length;
}
}
|