added extractReasoningFromText in extractCodeFromResult instead

This commit is contained in:
Andrew Pareles 2025-02-24 14:57:50 -08:00
parent aa9a102b71
commit 04efe1d235
3 changed files with 264 additions and 122 deletions

View file

@ -11,7 +11,6 @@ import { ResourceMap } from '../../../../../base/common/map.js';
import { isEqual } from '../../../../../base/common/resources.js';
import * as strings from '../../../../../base/common/strings.js';
import { URI } from '../../../../../base/common/uri.js';
import { stripThinkTags } from '../../common/chatModel.js';
import { IActiveCodeEditor, isCodeEditor, isDiffEditor } from '../../../../../editor/browser/editorBrowser.js';
import { IBulkEditService, ResourceTextEdit } from '../../../../../editor/browser/services/bulkEditService.js';
import { ICodeEditorService } from '../../../../../editor/browser/services/codeEditorService.js';
@ -383,11 +382,13 @@ function collectDocumentContextFromContext(context: ICodeBlockActionContext, res
}
function getChatConversation(context: ICodeBlockActionContext): (ConversationRequest | ConversationResponse)[] {
// TODO@aeschli for now create a conversation with just the current element
// this will be expanded in the future to include the request and any other responses
if (isResponseVM(context.element)) {
return [{
type: 'response',
// Strip think tags before computing diffs
message: stripThinkTags(context.element.response.toMarkdown()),
message: context.element.response.toMarkdown(),
references: getReferencesAsDocumentContext(context.element.contentReferences)
}];
} else if (isRequestVM(context.element)) {

View file

@ -8,7 +8,6 @@ import { DeferredPromise } from '../../../../base/common/async.js';
import { Emitter, Event } from '../../../../base/common/event.js';
import { IMarkdownString, MarkdownString, isMarkdownString } from '../../../../base/common/htmlContent.js';
import { Disposable } from '../../../../base/common/lifecycle.js';
import { SurroundingsRemover } from '../../../void/browser/helpers/extractCodeFromResult.js';
import { revive } from '../../../../base/common/marshalling.js';
import { equals } from '../../../../base/common/objects.js';
import { basename, isEqual } from '../../../../base/common/resources.js';
@ -250,11 +249,6 @@ export class Response extends Disposable implements IResponse {
updateContent(progress: IChatProgressResponseContent | IChatTextEdit | IChatTask, quiet?: boolean): void {
if (progress.kind === 'markdownContent') {
// Handle streaming for think tags
const remover = new ThinkTagSurroundingsRemover(progress.content.value);
const [delta, ignoredSuffix] = remover.deltaInfo(progress.content.value.length);
progress.content.value = delta;
const responsePartLength = this._responseParts.length - 1;
const lastResponsePart = this._responseParts[responsePartLength];
@ -360,78 +354,6 @@ export class Response extends Disposable implements IResponse {
}
}
/**
* Strips <think> tags and their content from a text string.
* Handles nested tags using a stack-based approach.
* @param text The text to strip tags from
* @returns The text with all <think> tags and their content removed
*/
export function stripThinkTags(text: string): string {
// Handle nested tags with a stack-based approach
let result = '';
let depth = 0;
let i = 0;
while (i < text.length) {
if (text.startsWith('<think>', i)) {
depth++;
i += 7; // length of '<think>'
} else if (text.startsWith('</think>', i)) {
if (depth > 0) depth--;
i += 8; // length of '</think>'
} else if (depth === 0) {
result += text[i];
i++;
} else {
i++;
}
}
return result;
}
class ThinkTagSurroundingsRemover extends SurroundingsRemover {
constructor(s: string) {
super(s);
}
removeThinkTags() {
// Handle token streaming at a more granular level
let foundTag = false;
// Try to remove opening tag, handling partial tokens
foundTag = this.removePrefix('<');
if (foundTag) {
foundTag = this.removePrefix('t');
if (foundTag) {
foundTag = this.removePrefix('h');
if (foundTag) {
foundTag = this.removePrefix('i');
if (foundTag) {
foundTag = this.removePrefix('n');
if (foundTag) {
foundTag = this.removePrefix('k');
if (foundTag) {
foundTag = this.removePrefix('>');
}
}
}
}
}
}
return foundTag;
}
deltaInfo(recentlyAddedTextLen: number) {
// Get the delta and suffix from parent class
const [delta, ignoredSuffix] = super.deltaInfo(recentlyAddedTextLen);
// Strip any think tags from the delta before returning
return [stripThinkTags(delta), ignoredSuffix] as const;
}
}
export class ChatResponseModel extends Disposable implements IChatResponseModel {
private readonly _onDidChange = this._register(new Emitter<void>());
readonly onDidChange = this._onDidChange.event;

View file

@ -3,57 +3,276 @@
* Licensed under the MIT License. See License.txt in the project root for license information.
*--------------------------------------------------------------------------------------------*/
import * as assert from 'assert';
import { Response, stripThinkTags } from '../../../common/chatModel';
import { MarkdownString } from '../../../../../../base/common/htmlContent';
import assert from 'assert';
import { timeout } from '../../../../../base/common/async.js';
import { MarkdownString } from '../../../../../base/common/htmlContent.js';
import { URI } from '../../../../../base/common/uri.js';
import { assertSnapshot } from '../../../../../base/test/common/snapshot.js';
import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../../base/test/common/utils.js';
import { OffsetRange } from '../../../../../editor/common/core/offsetRange.js';
import { Range } from '../../../../../editor/common/core/range.js';
import { IContextKeyService } from '../../../../../platform/contextkey/common/contextkey.js';
import { TestInstantiationService } from '../../../../../platform/instantiation/test/common/instantiationServiceMock.js';
import { MockContextKeyService } from '../../../../../platform/keybinding/test/common/mockKeybindingService.js';
import { ILogService, NullLogService } from '../../../../../platform/log/common/log.js';
import { IStorageService } from '../../../../../platform/storage/common/storage.js';
import { ChatAgentLocation, ChatAgentService, IChatAgentService } from '../../common/chatAgents.js';
import { ChatModel, ISerializableChatData1, ISerializableChatData2, ISerializableChatData3, normalizeSerializableChatData, Response } from '../../common/chatModel.js';
import { ChatRequestTextPart } from '../../common/chatParserTypes.js';
import { IExtensionService } from '../../../../services/extensions/common/extensions.js';
import { TestExtensionService, TestStorageService } from '../../../../test/common/workbenchTestServices.js';
suite('ChatModel - Think Tags', () => {
test('handles partial tokens during streaming', () => {
const response = new Response(new MarkdownString('<'));
assert.strictEqual(response.toString(), '');
suite('ChatModel', () => {
const testDisposables = ensureNoDisposablesAreLeakedInTestSuite();
response.updateContent({ kind: 'markdownContent', content: new MarkdownString('<t') });
assert.strictEqual(response.toString(), '');
let instantiationService: TestInstantiationService;
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(), '');
response.updateContent({ kind: 'markdownContent', content: new MarkdownString('<think>') });
assert.strictEqual(response.toString(), '');
response.updateContent({ kind: 'markdownContent', content: new MarkdownString('<think>test') });
assert.strictEqual(response.toString(), '');
response.updateContent({ kind: 'markdownContent', content: new MarkdownString('<think>test</think>') });
assert.strictEqual(response.toString(), 'test');
setup(async () => {
instantiationService = testDisposables.add(new TestInstantiationService());
instantiationService.stub(IStorageService, testDisposables.add(new TestStorageService()));
instantiationService.stub(ILogService, new NullLogService());
instantiationService.stub(IExtensionService, new TestExtensionService());
instantiationService.stub(IContextKeyService, new MockContextKeyService());
instantiationService.stub(IChatAgentService, instantiationService.createInstance(ChatAgentService));
});
test('handles malformed tags', () => {
assert.strictEqual(stripThinkTags('<think>unclosed'), '');
assert.strictEqual(stripThinkTags('</think>unopened'), 'unopened');
assert.strictEqual(stripThinkTags('<think>nested<think>tags</think></think>'), '');
test('Waits for initialization', async () => {
const model = testDisposables.add(instantiationService.createInstance(ChatModel, undefined, ChatAgentLocation.Panel));
let hasInitialized = false;
model.waitForInitialization().then(() => {
hasInitialized = true;
});
await timeout(0);
assert.strictEqual(hasInitialized, false);
model.startInitialize();
model.initialize(undefined);
await timeout(0);
assert.strictEqual(hasInitialized, true);
});
test('handles half-nested tags from conversation', () => {
const input = `<think>Okay, the user wants me to create nested <think> tags in my thought process. Let me start by recalling what they asked for. They initially mentioned using nested <think> tags with multiple layers. In my first attempt, I probably just used a single pair of <think> tags without nesting.
test('must call startInitialize before initialize', async () => {
const model = testDisposables.add(instantiationService.createInstance(ChatModel, undefined, ChatAgentLocation.Panel));
So, I need to make sure each opening <think> tag is properly closed with a </think> tag, and that these tags are nested. For example, starting with one <think> tag, then another inside it, and so on.</think>
let hasInitialized = false;
model.waitForInitialization().then(() => {
hasInitialized = true;
});
i tried`;
assert.strictEqual(stripThinkTags(input), 'i tried');
await timeout(0);
assert.strictEqual(hasInitialized, false);
assert.throws(() => model.initialize(undefined));
assert.strictEqual(hasInitialized, false);
});
test('preserves text mentions of tags', () => {
const input = 'Let me try with <think> tags and see how it works';
assert.strictEqual(stripThinkTags(input), 'Let me try with <think> tags and see how it works');
test('deinitialize/reinitialize', async () => {
const model = testDisposables.add(instantiationService.createInstance(ChatModel, undefined, ChatAgentLocation.Panel));
let hasInitialized = false;
model.waitForInitialization().then(() => {
hasInitialized = true;
});
model.startInitialize();
model.initialize(undefined);
await timeout(0);
assert.strictEqual(hasInitialized, true);
model.deinitialize();
let hasInitialized2 = false;
model.waitForInitialization().then(() => {
hasInitialized2 = true;
});
model.startInitialize();
model.initialize(undefined);
await timeout(0);
assert.strictEqual(hasInitialized2, true);
});
test('cannot initialize twice', async () => {
const model = testDisposables.add(instantiationService.createInstance(ChatModel, undefined, ChatAgentLocation.Panel));
model.startInitialize();
model.initialize(undefined);
assert.throws(() => model.initialize(undefined));
});
test('Initialization fails when model is disposed', async () => {
const model = testDisposables.add(instantiationService.createInstance(ChatModel, undefined, ChatAgentLocation.Panel));
model.dispose();
assert.throws(() => model.initialize(undefined));
});
test('removeRequest', async () => {
const model = testDisposables.add(instantiationService.createInstance(ChatModel, undefined, ChatAgentLocation.Panel));
model.startInitialize();
model.initialize(undefined);
const text = 'hello';
model.addRequest({ text, parts: [new ChatRequestTextPart(new OffsetRange(0, text.length), new Range(1, text.length, 1, text.length), text)] }, { variables: [] }, 0);
const requests = model.getRequests();
assert.strictEqual(requests.length, 1);
model.removeRequest(requests[0].id);
assert.strictEqual(model.getRequests().length, 0);
});
test('adoptRequest', async function () {
const model1 = testDisposables.add(instantiationService.createInstance(ChatModel, undefined, ChatAgentLocation.Editor));
const model2 = testDisposables.add(instantiationService.createInstance(ChatModel, undefined, ChatAgentLocation.Panel));
model1.startInitialize();
model1.initialize(undefined);
model2.startInitialize();
model2.initialize(undefined);
const text = 'hello';
const request1 = model1.addRequest({ text, parts: [new ChatRequestTextPart(new OffsetRange(0, text.length), new Range(1, text.length, 1, text.length), text)] }, { variables: [] }, 0);
assert.strictEqual(model1.getRequests().length, 1);
assert.strictEqual(model2.getRequests().length, 0);
assert.ok(request1.session === model1);
assert.ok(request1.response?.session === model1);
model2.adoptRequest(request1);
assert.strictEqual(model1.getRequests().length, 0);
assert.strictEqual(model2.getRequests().length, 1);
assert.ok(request1.session === model2);
assert.ok(request1.response?.session === model2);
model2.acceptResponseProgress(request1, { content: new MarkdownString('Hello'), kind: 'markdownContent' });
assert.strictEqual(request1.response.response.toString(), 'Hello');
});
});
suite('Response', () => {
const store = ensureNoDisposablesAreLeakedInTestSuite();
test('mergeable markdown', async () => {
const response = store.add(new Response([]));
response.updateContent({ content: new MarkdownString('markdown1'), kind: 'markdownContent' });
response.updateContent({ content: new MarkdownString('markdown2'), kind: 'markdownContent' });
await assertSnapshot(response.value);
assert.strictEqual(response.toString(), 'markdown1markdown2');
});
test('not mergeable markdown', async () => {
const response = store.add(new Response([]));
const md1 = new MarkdownString('markdown1');
md1.supportHtml = true;
response.updateContent({ content: md1, kind: 'markdownContent' });
response.updateContent({ content: new MarkdownString('markdown2'), kind: 'markdownContent' });
await assertSnapshot(response.value);
});
test('inline reference', async () => {
const response = store.add(new Response([]));
response.updateContent({ content: new MarkdownString('text before'), kind: 'markdownContent' });
response.updateContent({ inlineReference: URI.parse('https://microsoft.com'), kind: 'inlineReference' });
response.updateContent({ content: new MarkdownString('text after'), kind: 'markdownContent' });
await assertSnapshot(response.value);
});
});
suite('normalizeSerializableChatData', () => {
ensureNoDisposablesAreLeakedInTestSuite();
test('v1', () => {
const v1Data: ISerializableChatData1 = {
creationDate: Date.now(),
initialLocation: undefined,
isImported: false,
requesterAvatarIconUri: undefined,
requesterUsername: 'me',
requests: [],
responderAvatarIconUri: undefined,
responderUsername: 'bot',
sessionId: 'session1',
};
const newData = normalizeSerializableChatData(v1Data);
assert.strictEqual(newData.creationDate, v1Data.creationDate);
assert.strictEqual(newData.lastMessageDate, v1Data.creationDate);
assert.strictEqual(newData.version, 3);
assert.ok('customTitle' in newData);
});
test('v2', () => {
const v2Data: ISerializableChatData2 = {
version: 2,
creationDate: 100,
lastMessageDate: Date.now(),
initialLocation: undefined,
isImported: false,
requesterAvatarIconUri: undefined,
requesterUsername: 'me',
requests: [],
responderAvatarIconUri: undefined,
responderUsername: 'bot',
sessionId: 'session1',
computedTitle: 'computed title'
};
const newData = normalizeSerializableChatData(v2Data);
assert.strictEqual(newData.version, 3);
assert.strictEqual(newData.creationDate, v2Data.creationDate);
assert.strictEqual(newData.lastMessageDate, v2Data.lastMessageDate);
assert.strictEqual(newData.customTitle, v2Data.computedTitle);
});
test('old bad data', () => {
const v1Data: ISerializableChatData1 = {
// Testing the scenario where these are missing
sessionId: undefined!,
creationDate: undefined!,
initialLocation: undefined,
isImported: false,
requesterAvatarIconUri: undefined,
requesterUsername: 'me',
requests: [],
responderAvatarIconUri: undefined,
responderUsername: 'bot',
};
const newData = normalizeSerializableChatData(v1Data);
assert.strictEqual(newData.version, 3);
assert.ok(newData.creationDate > 0);
assert.ok(newData.lastMessageDate > 0);
assert.ok(newData.sessionId);
});
test('v3 with bug', () => {
const v3Data: ISerializableChatData3 = {
// Test case where old data was wrongly normalized and these fields were missing
creationDate: undefined!,
lastMessageDate: undefined!,
version: 3,
initialLocation: undefined,
isImported: false,
requesterAvatarIconUri: undefined,
requesterUsername: 'me',
requests: [],
responderAvatarIconUri: undefined,
responderUsername: 'bot',
sessionId: 'session1',
customTitle: 'computed title'
};
const newData = normalizeSerializableChatData(v3Data);
assert.strictEqual(newData.version, 3);
assert.ok(newData.creationDate > 0);
assert.ok(newData.lastMessageDate > 0);
assert.ok(newData.sessionId);
});
});