From d74fc714d438cc680dcdc20b7a32ae25141e2a1f Mon Sep 17 00:00:00 2001 From: Bernhard Jackiewicz Date: Fri, 2 Oct 2026 19:36:13 +0200 Subject: [PATCH] fix(ui): count an answer that ends on its last allowed token as finished The worker inferred "length" whenever the answer used every allowed token. Transformers.js also hands the end-of-sequence token to the streamer, so an answer that ended exactly on the limit was reported as cut and the chat showed "Shortened (token limit)". The worker now remembers the last token and reports "length" only when it is not an end-of-sequence id of the model or tokenizer. Signed-off-by: Bernhard Jackiewicz --- graph-ui/src/browser-ai/browser-ai.worker.test.ts | 13 +++++++++---- graph-ui/src/browser-ai/browser-ai.worker.ts | 11 +++++++---- 2 files changed, 16 insertions(+), 8 deletions(-) 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) => {