Improve token streaming support with character-by-character matching

Co-Authored-By: Jack Hacksman <slack@hannis.io>
This commit is contained in:
Devin AI 2025-02-16 05:06:37 +00:00
parent b24cd1ecb7
commit 2f55580e67
2 changed files with 18 additions and 22 deletions

View file

@ -371,7 +371,6 @@ export function stripThinkTags(text: string): string {
let result = '';
let depth = 0;
let i = 0;
let inPartialTag = false;
while (i < text.length) {
if (text.startsWith('<think>', i)) {
@ -380,24 +379,12 @@ export function stripThinkTags(text: string): string {
} else if (text.startsWith('</think>', i)) {
if (depth > 0) depth--;
i += 8; // length of '</think>'
} else if (text.startsWith('<thi', i)) {
// Handle partial opening tag during streaming
inPartialTag = true;
i += 4;
} else if (text.startsWith('</thi', i)) {
// Handle partial closing tag during streaming
inPartialTag = true;
i += 5;
} else if (depth === 0 && !inPartialTag) {
} else if (depth === 0) {
result += text[i];
i++;
} else {
i++;
}
// Reset partial tag flag after moving past potential tag
if (inPartialTag && !text.startsWith('nk>', i)) {
inPartialTag = false;
}
}
return result;
@ -409,15 +396,18 @@ class ThinkTagSurroundingsRemover extends SurroundingsRemover {
}
removeThinkTags() {
// Try to remove opening tag, handling partial tokens during streaming
const foundOpenTag = this.removePrefix('<think>');
if (!foundOpenTag) {
// Let removePrefix handle partial matches character by character
this.removePrefix('<');
this.removePrefix('think');
this.removePrefix('>');
// Handle token streaming character by character
const chars = ['<', 't', 'h', 'i', 'n', 'k', '>'];
let foundTag = true;
for (const char of chars) {
if (!this.removePrefix(char)) {
foundTag = false;
break;
}
}
return foundOpenTag;
return foundTag;
}
deltaInfo(recentlyAddedTextLen: number) {

View file

@ -15,9 +15,15 @@ suite('ChatModel - Think Tags', () => {
response.updateContent({ kind: 'markdownContent', content: new MarkdownString('<t') });
assert.strictEqual(response.toString(), '');
response.updateContent({ kind: 'markdownContent', content: new MarkdownString('<th') });
assert.strictEqual(response.toString(), '');
response.updateContent({ kind: 'markdownContent', content: new MarkdownString('<thi') });
assert.strictEqual(response.toString(), '');
response.updateContent({ kind: 'markdownContent', content: new MarkdownString('<thin') });
assert.strictEqual(response.toString(), '');
response.updateContent({ kind: 'markdownContent', content: new MarkdownString('<think') });
assert.strictEqual(response.toString(), '');