diff --git a/graph-ui/src/browser-ai/browser-ai.worker.test.ts b/graph-ui/src/browser-ai/browser-ai.worker.test.ts index 0c3678be5..1c42460f4 100644 --- a/graph-ui/src/browser-ai/browser-ai.worker.test.ts +++ b/graph-ui/src/browser-ai/browser-ai.worker.test.ts @@ -22,8 +22,8 @@ vi.mock('@huggingface/transformers', () => ({ }, TextStreamer: class { callback: (chunk: string) => void; - token: () => void; - constructor(_tokenizer: unknown, options: { callback_function: (chunk: string) => void; token_callback_function: () => void }) { + token: (tokens?: bigint[]) => void; + constructor(_tokenizer: unknown, options: { callback_function: (chunk: string) => void; token_callback_function: (tokens?: bigint[]) => void }) { this.callback = options.callback_function; this.token = options.token_callback_function; } }, @@ -133,10 +133,11 @@ describe('browser generation worker', () => { }); it('reports whether an answer ended at the output limit, at the end of sequence, or on Stop', async () => { + fake.modelLoad.mockResolvedValue({ generate: fake.generate, generation_config: { eos_token_id: [2, 3] } }); await send({ id: 1, kind: 'prepare' }); const messages = [{ role: 'user' as const, content: 'List' }]; - const emit = (tokens: number) => async ({ streamer }: { streamer: { callback: (chunk: string) => void; token: () => void } }) => { - for (let index = 0; index < tokens; index++) { streamer.token(); streamer.callback('x'); } + const emit = (tokens: number, last = 7n) => async ({ streamer }: { streamer: { callback: (chunk: string) => void; token: (tokens?: bigint[]) => void } }) => { + for (let index = 0; index < tokens; index++) { streamer.token([index + 1 === tokens ? last : 9n]); streamer.callback('x'); } }; fake.generate.mockImplementationOnce(emit(64)); await send({ id: 2, kind: 'chat', messages, maxOutputTokens: 64 }); @@ -144,6 +145,10 @@ describe('browser generation worker', () => { fake.generate.mockImplementationOnce(emit(10)); await send({ id: 3, kind: 'chat', messages, maxOutputTokens: 64 }); expect(replies().at(-1)).toEqual({ id: 3, kind: 'answer', output: 'x'.repeat(10), stopReason: 'eos' }); + // The streamer also receives the end-of-sequence token; ending on the last allowed token is not a cut. + fake.generate.mockImplementationOnce(emit(64, 3n)); + await send({ id: 4, kind: 'chat', messages, maxOutputTokens: 64 }); + expect(replies().at(-1)).toEqual({ id: 4, kind: 'answer', output: 'x'.repeat(64), stopReason: 'eos' }); }); it('rejects unknown profiles and oversized automatic answers without invalidating the model', async () => { diff --git a/graph-ui/src/browser-ai/browser-ai.worker.ts b/graph-ui/src/browser-ai/browser-ai.worker.ts index 234f296c4..4fbf5b2df 100644 --- a/graph-ui/src/browser-ai/browser-ai.worker.ts +++ b/graph-ui/src/browser-ai/browser-ai.worker.ts @@ -158,11 +158,11 @@ async function generate(id: number, messages: readonly BrowserChatMessage[], opt throw new Error(`This conversation uses ${count.toLocaleString()} input tokens. The browser limit is ${selected.contextTokens.toLocaleString()}, including ${outputTokens} reserved for the answer. Remove earlier messages or attach a smaller selection; no code was truncated.`); } stopping.reset(); - let answer = '', generated = 0; + let answer = '', generated = 0, last: bigint | undefined; const streamer = new TextStreamer(tokenizer!, { skip_prompt: true, callback_function: chunk => { answer += chunk; post({ id, kind: 'token', output: chunk }); }, - token_callback_function: () => { generated += 1; }, + token_callback_function: tokens => { generated += 1; last = tokens?.at(-1); }, }); const controls = REPETITION_CONTROLS[automaticExplanation ? 'automatic-explanation' : 'chat']; const logitsProcessor = new LogitsProcessorList(); @@ -181,8 +181,11 @@ async function generate(id: number, messages: readonly BrowserChatMessage[], opt } if (fatalFailure) throw fatalFailure; if (!answer.trim() && !stopping.interrupted) throw new Error('The model returned no answer.'); - // A stopped answer is reported as stopped; a full output budget means the answer was cut. - return { output: answer, stopReason: stopping.interrupted ? 'interrupted' : generated >= outputTokens ? 'length' : 'eos' }; + // A stopped answer is reported as stopped; a full output budget means the answer was cut, + // unless its last token was the end of sequence (the streamer also receives that token). + const ends = [model!.generation_config?.eos_token_id ?? [], tokenizer!.eos_token_id ?? []].flat().map(token => BigInt(token)); + const ended = last !== undefined && ends.includes(last); + return { output: answer, stopReason: stopping.interrupted ? 'interrupted' : generated >= outputTokens && !ended ? 'length' : 'eos' }; } self.onmessage = async (event: MessageEvent) => {