diff --git a/src/backend/drivers/ai-chat/providers/openrouter/OpenRouterProvider.test.ts b/src/backend/drivers/ai-chat/providers/openrouter/OpenRouterProvider.test.ts index fb5b67301..bc2b2d0b5 100644 --- a/src/backend/drivers/ai-chat/providers/openrouter/OpenRouterProvider.test.ts +++ b/src/backend/drivers/ai-chat/providers/openrouter/OpenRouterProvider.test.ts @@ -354,11 +354,36 @@ describe('OpenRouterProvider.complete request shape', () => { }); }); - it('retries without max_tokens when OpenRouter rejects with a context-length error', async () => { + it('retries under the room left by the window when OpenRouter reports it', async () => { + const { provider } = makeProvider(); + + const ctxErr = { + error: { + message: + "This endpoint's maximum context length is 4096 tokens. However, you requested 5000 tokens (900 of text input, 4100 in the output).", + }, + }; + createMock + .mockRejectedValueOnce(ctxErr) + .mockResolvedValueOnce(baseCompletion); + + await withTestActor(() => + provider.complete({ + model: 'openrouter:openai/gpt-5-nano', + messages: [{ role: 'user', content: 'hi' }], + max_tokens: 4100, + }), + ); + + expect(createMock).toHaveBeenCalledTimes(2); + // 4096 of window less the 900 input tokens the rejection reported. + expect(createMock.mock.calls[1]![0].max_tokens).toBe(3196); + expect(createMock.mock.calls[0]![0].max_tokens).toBe(4100); + }); + + it('sizes the retry from the prompt estimate when the rejection omits token counts', async () => { const { provider } = makeProvider(); - // First call: simulate the OpenRouter "context length" rejection - // shape the provider catches by message prefix. const ctxErr = { error: { message: @@ -372,17 +397,16 @@ describe('OpenRouterProvider.complete request shape', () => { await withTestActor(() => provider.complete({ model: 'openrouter:openai/gpt-5-nano', - messages: [{ role: 'user', content: 'hi' }], + // ~1000 estimated tokens, doubled: the retry assumes the + // estimator's whitespace-poor worst case. + messages: [{ role: 'user', content: 'x'.repeat(8000) }], max_tokens: 9999999, }), ); - // Provider mutates a single completionParams object across both calls - // (`delete completionParams.max_tokens` after the first throw), so we - // can only assert that two calls happened and the surviving shape no - // longer carries max_tokens. expect(createMock).toHaveBeenCalledTimes(2); - expect('max_tokens' in createMock.mock.calls[1]![0]).toBe(false); + // 4096 of window less the 2000 the prompt is assumed to occupy. + expect(createMock.mock.calls[1]![0].max_tokens).toBe(2096); }); it('rethrows non-context-length errors without retrying', async () => { @@ -407,6 +431,27 @@ describe('OpenRouterProvider.complete request shape', () => { expect(recordSpy).not.toHaveBeenCalled(); expect(logSpy).toHaveBeenCalled(); }); + + it('rethrows a transport error as itself rather than masking it', async () => { + const { provider } = makeProvider(); + // No `.error` on the object: the shape a socket failure arrives in. + const transportError = new Error('socket hang up'); + createMock.mockRejectedValueOnce(transportError); + const logSpy = vi.spyOn(console, 'log').mockImplementation(() => {}); + + await expect( + withTestActor(() => + provider.complete({ + model: 'openrouter:openai/gpt-5-nano', + messages: [{ role: 'user', content: 'boom' }], + max_tokens: 100, + }), + ), + ).rejects.toBe(transportError); + + expect(createMock).toHaveBeenCalledTimes(1); + expect(logSpy).toHaveBeenCalled(); + }); }); // ── Non-stream completion: cost calculator branches ───────────────── diff --git a/src/backend/drivers/ai-chat/providers/openrouter/OpenRouterProvider.ts b/src/backend/drivers/ai-chat/providers/openrouter/OpenRouterProvider.ts index c673cb600..2170d6e95 100644 --- a/src/backend/drivers/ai-chat/providers/openrouter/OpenRouterProvider.ts +++ b/src/backend/drivers/ai-chat/providers/openrouter/OpenRouterProvider.ts @@ -25,6 +25,10 @@ import { Context } from '../../../../core/context.js'; import type { MeteringService } from '../../../../services/metering/MeteringService.js'; import { kv } from '../../../../util/kvSingleton.js'; import * as OpenAIUtil from '../../utils/OpenAIUtil.js'; +import { + contextLengthRetryParams, + isContextLengthError, +} from '../../utils/contextLimit.js'; import type { IChatModel, IChatProvider, @@ -146,27 +150,23 @@ export class OpenRouterProvider implements IChatProvider { completion = await this.#openai.chat.completions.create(completionParams); } catch (e: unknown) { - // If you overestimate allowed max_tokens on openrouter then it will throw an error. - // Since we know the user has enough for the query anyways, we should reexecute the - // request without max_tokens. - const err = e as { error: Error }; - if ( - err && - err.error && - err.error.message && - err.error.message.startsWith( - "This endpoint's maximum context length is ", - ) - ) { - delete completionParams.max_tokens; - completion = - await this.#openai.chat.completions.create( - completionParams, - ); - } else { - console.log('Openarouter error: ', err.error.message); + if (!isContextLengthError(e)) { + console.log( + 'Openrouter error: ', + (e as { error?: { message?: string } })?.error?.message, + ); throw e; } + // OpenRouter rejects an overlarge max_tokens rather than + // truncating. Retry under the room the window leaves, still + // bounded by the cap the credit gate set. + const retryParams = contextLengthRetryParams(completionParams, { + error: e, + contextWindow: modelUsed.context, + }); + if (!retryParams) throw e; + completion = + await this.#openai.chat.completions.create(retryParams); } return OpenAIUtil.handle_completion_output({ diff --git a/src/backend/drivers/ai-chat/providers/together/TogetherAIProvider.test.ts b/src/backend/drivers/ai-chat/providers/together/TogetherAIProvider.test.ts index bcac330a4..2c3393f62 100644 --- a/src/backend/drivers/ai-chat/providers/together/TogetherAIProvider.test.ts +++ b/src/backend/drivers/ai-chat/providers/together/TogetherAIProvider.test.ts @@ -215,7 +215,8 @@ describe('TogetherAIProvider model catalog', () => { (m) => m.id === 'togetherai:Qwen/Qwen2.5-7B-Instruct-Turbo', )!; // The advertised context window is unchanged; only the output cap - // leaves room for the driver's under-counting input estimator. + // leaves room for an input estimate that can run half low on + // whitespace-poor prompts. expect(model.context).toBe(32768); expect(model.max_tokens).toBe(Math.floor(32768 * 0.95)); }); @@ -386,29 +387,152 @@ describe('TogetherAIProvider.complete request shape', () => { }, }; - it('retries without max_tokens when Together rejects with a context-length error', async () => { + it('retries under the room left by the window when Together rejects with a context-length error', async () => { const { provider } = makeProvider(); createMock .mockRejectedValueOnce(contextLengthError) .mockResolvedValueOnce(baseCompletion); + const tools = [ + { + type: 'function', + function: { name: 'get_weather', parameters: {} }, + }, + ]; await withTestActor(() => provider.complete({ model: 'togetherai:Qwen/Qwen2.5-7B-Instruct-Turbo', messages: [{ role: 'user', content: 'hi' }], + tools, max_tokens: 1048573, + } as never), + ); + + expect(createMock).toHaveBeenCalledTimes(2); + // 1048576 of window less the 11 input tokens the rejection reported. + expect(createMock.mock.calls[1]![0].max_tokens).toBe(1048565); + // Resizing the request must not drop what it was carrying. + expect(createMock.mock.calls[1]![0].tools).toHaveLength(1); + // The first attempt's params are left as they were sent. + expect(createMock.mock.calls[0]![0].max_tokens).toBe(1048573); + // The retry is still metered exactly once. + expect(recordSpy).toHaveBeenCalledTimes(1); + }); + + it('gives up when the estimate leaves no smaller cap to retry with', async () => { + const { provider } = makeProvider(); + + // The rejection carries no counts, the prompt estimates at ~100 + // tokens, and 3990 already sits under the 3996 that leaves — so the + // only retry available is the request just rejected. + const noCounts = { + status: 400, + error: { + error: { + message: + "This model's maximum context length is 4096 tokens.", + }, + }, + }; + createMock.mockRejectedValueOnce(noCounts); + + await expect( + withTestActor(() => + provider.complete({ + model: 'togetherai:Qwen/Qwen2.5-7B-Instruct-Turbo', + messages: [{ role: 'user', content: 'x'.repeat(400) }], + max_tokens: 3990, + }), + ), + ).rejects.toBe(noCounts); + + expect(createMock).toHaveBeenCalledTimes(1); + }); + + it('caps a retry that had no cap to begin with', async () => { + const { provider } = makeProvider(); + + createMock + .mockRejectedValueOnce(contextLengthError) + .mockResolvedValueOnce(baseCompletion); + + // Output priced at zero leaves the gate nothing to cap, so the first + // attempt goes out without max_tokens and the retry sizes to the room. + await withTestActor(() => + provider.complete({ + model: 'togetherai:Qwen/Qwen2.5-7B-Instruct-Turbo', + messages: [{ role: 'user', content: 'hi' }], }), ); - // Provider mutates a single completionParams object across both calls - // (`delete completionParams.max_tokens` after the first throw), so the - // first call's recorded args are retroactively altered — only the - // surviving shape is worth asserting on. expect(createMock).toHaveBeenCalledTimes(2); - expect('max_tokens' in createMock.mock.calls[1]![0]).toBe(false); - // The retry is still metered exactly once. - expect(recordSpy).toHaveBeenCalledTimes(1); + expect('max_tokens' in createMock.mock.calls[0]![0]).toBe(false); + expect(createMock.mock.calls[1]![0].max_tokens).toBe(1048565); + }); + + it('gives up instead of retrying when the prompt alone fills the window', async () => { + const { provider } = makeProvider(); + + const noRoom = { + status: 400, + error: { + error: { + message: + "This model's maximum context length is 8192 tokens. However, your messages resulted in 9000 tokens.", + }, + }, + }; + createMock.mockRejectedValueOnce(noRoom); + + await expect( + withTestActor(() => + provider.complete({ + model: 'togetherai:Qwen/Qwen2.5-7B-Instruct-Turbo', + messages: [{ role: 'user', content: 'hi' }], + max_tokens: 500, + }), + ), + ).rejects.toBe(noRoom); + + expect(createMock).toHaveBeenCalledTimes(1); + }); + + it('gives up rather than retrying uncapped when no window can be determined', async () => { + // A listing with no context_length leaves the model without a + // declared window, and this rejection carries no figure either. + modelsListMock.mockResolvedValue([ + { + id: 'Qwen/Qwen2.5-7B-Instruct-Turbo', + type: 'chat', + display_name: 'Qwen 2.5 7B Instruct Turbo', + pricing: { input: 20, output: 20 }, + }, + ]); + const { provider } = makeProvider(); + + const noWindow = { + status: 400, + error: { + error: { + message: + 'Request exceeds the maximum context length for this model.', + }, + }, + }; + createMock.mockRejectedValueOnce(noWindow); + + await expect( + withTestActor(() => + provider.complete({ + model: 'togetherai:Qwen/Qwen2.5-7B-Instruct-Turbo', + messages: [{ role: 'user', content: 'hi' }], + max_tokens: 500, + }), + ), + ).rejects.toBe(noWindow); + + expect(createMock).toHaveBeenCalledTimes(1); }); it('retries a streaming request the same way', async () => { @@ -428,7 +552,7 @@ describe('TogetherAIProvider.complete request shape', () => { ); expect(createMock).toHaveBeenCalledTimes(2); - expect('max_tokens' in createMock.mock.calls[1]![0]).toBe(false); + expect(createMock.mock.calls[1]![0].max_tokens).toBe(1048565); expect(createMock.mock.calls[1]![0].stream).toBe(true); }); diff --git a/src/backend/drivers/ai-chat/providers/together/TogetherAIProvider.ts b/src/backend/drivers/ai-chat/providers/together/TogetherAIProvider.ts index 8cc9603a2..6bbd45337 100644 --- a/src/backend/drivers/ai-chat/providers/together/TogetherAIProvider.ts +++ b/src/backend/drivers/ai-chat/providers/together/TogetherAIProvider.ts @@ -23,6 +23,10 @@ import type { MeteringService } from '../../../../services/metering/MeteringServ import { kv } from '../../../../util/kvSingleton.js'; import { IChatModel, IChatProvider, ICompleteArguments } from '../../types.js'; import * as OpenAIUtil from '../../utils/OpenAIUtil.js'; +import { + contextLengthRetryParams, + isContextLengthError, +} from '../../utils/contextLimit.js'; import { modelLookupNames } from '../../utils/modelRouting.js'; const TOGETHER_AI_CHAT_COST_MAP: Record = { @@ -30,24 +34,6 @@ const TOGETHER_AI_CHAT_COST_MAP: Record = { completion_tokens: 'output', }; -/** - * Whether the SDK rejected a request because the prompt plus the requested - * output exceeds the model's context window. Unlike the OpenAI SDK, Together's - * `APIError.error` is the whole response body, so the provider message sits one - * level deeper; `message` is the stringified body and covers older shapes. - */ -const isContextLengthError = (e: unknown) => { - const err = e as { - error?: { error?: { message?: string } }; - message?: string; - }; - const message = err?.error?.error?.message ?? err?.message; - return ( - typeof message === 'string' && - message.includes('maximum context length') - ); -}; - export class TogetherAIProvider implements IChatProvider { #together: Together; @@ -103,9 +89,9 @@ export class TogetherAIProvider implements IChatProvider { ), }, // Together only reports a context length. The driver caps - // output at max_tokens minus an estimated input count, and - // that estimate runs low — reserve headroom so a short - // prompt doesn't ask for more than the context allows. + // output at max_tokens minus an estimated input count, which + // runs low on whitespace-poor prompts — reserve headroom so + // the cap doesn't overshoot the context as often. max_tokens: model.context_length ? Math.floor(model.context_length * 0.95) : 8000, @@ -177,13 +163,17 @@ export class TogetherAIProvider implements IChatProvider { completion = await this.#together.chat.completions.create(completionParams); } catch (e: unknown) { - // An overestimated max_tokens makes Together reject the request - // outright rather than truncating. The user can afford the query - // either way, so retry once without the cap. + // Together rejects an overlarge max_tokens outright rather than + // truncating. Retry under the room the window leaves, still + // bounded by the cap the credit gate set. if (!isContextLengthError(e)) throw e; - delete completionParams.max_tokens; + const retryParams = contextLengthRetryParams(completionParams, { + error: e, + contextWindow: modelUsed.context, + }); + if (!retryParams) throw e; completion = - await this.#together.chat.completions.create(completionParams); + await this.#together.chat.completions.create(retryParams); } return OpenAIUtil.handle_completion_output({ diff --git a/src/backend/drivers/ai-chat/utils/contextLimit.test.ts b/src/backend/drivers/ai-chat/utils/contextLimit.test.ts new file mode 100644 index 000000000..947b10b1d --- /dev/null +++ b/src/backend/drivers/ai-chat/utils/contextLimit.test.ts @@ -0,0 +1,245 @@ +/* + * Copyright (C) 2024-present Puter Technologies Inc. + * + * This file is part of Puter. + * + * Puter is free software: you can redistribute it and/or modify + * it under the terms of the GNU Affero General Public License as published + * by the Free Software Foundation, either version 3 of the License, or + * (at your option) any later version. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * GNU Affero General Public License for more details. + * + * You should have received a copy of the GNU Affero General Public License + * along with this program. If not, see . + */ + +import { describe, expect, it } from 'vitest'; +import { + contextLengthRetryCap, + contextLengthRetryParams, + isContextLengthError, +} from './contextLimit.js'; + +// Together's APIError carries the whole response body on `.error`, so the +// provider message sits one level deeper than the OpenAI SDK puts it. +const togetherError = (message: string) => ({ + status: 400, + message: `400 {"error":{"message":${JSON.stringify(message)}}}`, + error: { error: { message, type: 'invalid_request_error' } }, +}); + +const openRouterError = (message: string) => ({ error: { message } }); + +const TOGETHER_MESSAGE = + "Failed to start generation: The input token count (11) plus the requested output count (1048573) exceeds the model's maximum context length (1048576)"; + +const OPENROUTER_MESSAGE = + "This endpoint's maximum context length is 4096 tokens. However, you requested 5000 tokens (900 of text input, 4100 in the output)."; + +describe('isContextLengthError', () => { + it('recognises both vendors regardless of where the SDK nests the message', () => { + expect(isContextLengthError(togetherError(TOGETHER_MESSAGE))).toBe( + true, + ); + expect(isContextLengthError(openRouterError(OPENROUTER_MESSAGE))).toBe( + true, + ); + }); + + it('accepts the phrase wherever a vendor puts it, in any case', () => { + expect( + isContextLengthError( + new Error("This model's maximum context length is 4096 tokens"), + ), + ).toBe(true); + expect( + isContextLengthError( + openRouterError('MAXIMUM CONTEXT LENGTH exceeded for this request'), + ), + ).toBe(true); + }); + + it('ignores unrelated failures and non-object throws', () => { + expect(isContextLengthError(openRouterError('Some other failure'))).toBe( + false, + ); + expect(isContextLengthError(new Error('socket hang up'))).toBe(false); + expect(isContextLengthError(undefined)).toBe(false); + expect(isContextLengthError('boom')).toBe(false); + }); +}); + +describe('contextLengthRetryCap', () => { + it('sizes the retry to the room the window leaves', () => { + expect( + contextLengthRetryCap({ + error: togetherError(TOGETHER_MESSAGE), + cap: 1048573, + contextWindow: undefined, + }), + ).toBe(1048565); + + expect( + contextLengthRetryCap({ + error: openRouterError(OPENROUTER_MESSAGE), + cap: 4100, + contextWindow: undefined, + }), + ).toBe(3196); + }); + + it('sizes to the room alone when the gate set no cap', () => { + expect( + contextLengthRetryCap({ + error: togetherError(TOGETHER_MESSAGE), + cap: undefined, + contextWindow: undefined, + }), + ).toBe(1048565); + }); + + it('never returns more than the cap the credit gate set', () => { + expect( + contextLengthRetryCap({ + error: togetherError(TOGETHER_MESSAGE), + cap: 64, + contextWindow: undefined, + }), + ).toBe(64); + }); + + it('falls back to the declared window when the message omits it', () => { + expect( + contextLengthRetryCap({ + error: openRouterError( + "This endpoint's maximum context length is exceeded. However, you requested 5000 tokens (900 of text input, 4100 in the output).", + ), + cap: 4100, + contextWindow: 4096, + }), + ).toBe(3196); + }); + + it('reports no room when the prompt alone fills the window', () => { + const cap = contextLengthRetryCap({ + error: openRouterError( + "This model's maximum context length is 8192 tokens. However, your messages resulted in 9000 tokens.", + ), + cap: 500, + contextWindow: undefined, + }); + expect(cap).toBeLessThan(1); + }); + + it('measures the prompt when the rejection omits the input count', () => { + expect( + contextLengthRetryCap({ + error: openRouterError( + "This endpoint's maximum context length is 4096 tokens.", + ), + cap: 9999999, + contextWindow: undefined, + request: { messages: [{ role: 'user', content: 'x'.repeat(8000) }] }, + }), + ).toBe(2096); + }); + + it('returns undefined when neither the message nor a request is available', () => { + expect( + contextLengthRetryCap({ + error: openRouterError( + "This endpoint's maximum context length is 4096 tokens.", + ), + cap: 9999999, + contextWindow: 4096, + }), + ).toBeUndefined(); + }); +}); + +describe('contextLengthRetryParams', () => { + it('clamps max_tokens instead of dropping it, leaving the original intact', () => { + const params = { model: 'm', max_tokens: 1048573 }; + const retry = contextLengthRetryParams(params, { + error: togetherError(TOGETHER_MESSAGE), + contextWindow: undefined, + }); + + expect(retry).toEqual({ model: 'm', max_tokens: 1048565 }); + expect(params.max_tokens).toBe(1048573); + }); + + it('keeps a cap whenever a window is known, even with no reported counts', () => { + const retry = contextLengthRetryParams( + { + model: 'm', + max_tokens: 9999999, + messages: [{ role: 'user', content: 'x'.repeat(4000) }], + }, + { + error: openRouterError( + "This endpoint's maximum context length is 4096 tokens.", + ), + contextWindow: 4096, + }, + ); + + // 4096 of window less the 1000 the prompt is assumed to occupy. + expect(retry!.max_tokens).toBe(3096); + }); + + it('gives up rather than retrying uncapped when no window can be determined', () => { + expect( + contextLengthRetryParams( + { model: 'm', max_tokens: 9999999 }, + { + // Carries the phrase but no window figure, and the model + // declares none either. + error: openRouterError( + 'Request exceeds the maximum context length for this model.', + ), + contextWindow: undefined, + }, + ), + ).toBeUndefined(); + }); + + it('gives up when the estimate leaves no smaller cap to retry with', () => { + // No counts in the rejection, a prompt the estimator sizes at ~100 + // tokens, and a cap already under the room that leaves: the retry + // would be the request just rejected. + expect( + contextLengthRetryParams( + { + model: 'm', + max_tokens: 3990, + messages: [{ role: 'user', content: 'x'.repeat(400) }], + }, + { + error: openRouterError( + "This endpoint's maximum context length is 4096 tokens.", + ), + contextWindow: undefined, + }, + ), + ).toBeUndefined(); + }); + + it('gives up rather than retrying when no output can fit', () => { + expect( + contextLengthRetryParams( + { model: 'm', max_tokens: 500 }, + { + error: openRouterError( + "This model's maximum context length is 8192 tokens. However, your messages resulted in 9000 tokens.", + ), + contextWindow: undefined, + }, + ), + ).toBeUndefined(); + }); +}); diff --git a/src/backend/drivers/ai-chat/utils/contextLimit.ts b/src/backend/drivers/ai-chat/utils/contextLimit.ts new file mode 100644 index 000000000..5643b96fc --- /dev/null +++ b/src/backend/drivers/ai-chat/utils/contextLimit.ts @@ -0,0 +1,147 @@ +/* + * Copyright (C) 2024-present Puter Technologies Inc. + * + * This file is part of Puter. + * + * Puter is free software: you can redistribute it and/or modify + * it under the terms of the GNU Affero General Public License as published + * by the Free Software Foundation, either version 3 of the License, or + * (at your option) any later version. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * GNU Affero General Public License for more details. + * + * You should have received a copy of the GNU Affero General Public License + * along with this program. If not, see . + */ + +/** + * Recovery from an upstream rejection for overflowing a model's context window. + * + * Providers reject rather than truncate when the prompt plus the requested + * output doesn't fit, so the request is retried under the room the window + * actually leaves — never above the cap the credit gate set, when it set one, + * which is what keeps the attempt bounded by the caller's balance. Output + * priced at zero leaves no cap, and the retry then carries `window - input` + * alone. Fitting the window is best effort; staying under that cap is not. + */ + +import { estimatePromptTokens } from './usageEstimate.js'; + +/** + * How far to over-count an estimated prompt. The shared estimator runs about + * half low on JSON and code and lower still on denser text, and a retry sized + * above the window earns the same rejection a second time — so assume the + * common worst case rather than the central one. + */ +const ESTIMATED_PROMPT_MARGIN = 2; + +/** Vendor phrasings for the context window, tried in order. */ +const CONTEXT_WINDOW_PATTERNS = [ + /maximum context length \((\d+)\)/i, + /maximum context length is (\d+)/i, +]; + +/** Vendor phrasings for the prompt's measured token count. */ +const INPUT_TOKEN_PATTERNS = [ + /input token count \((\d+)\)/i, + /(\d+) of text input/i, + /resulted in (\d+) tokens/i, +]; + +/** + * The upstream message, wherever the SDK put it. Together's `error` is the + * whole response body, so its message sits one level deeper than OpenAI's. + */ +const errorMessage = (e: unknown): string | undefined => { + const err = e as { + error?: { message?: string; error?: { message?: string } }; + message?: string; + }; + const message = + err?.error?.error?.message ?? err?.error?.message ?? err?.message; + return typeof message === 'string' ? message : undefined; +}; + +/** Whether the upstream rejected the request for exceeding its context window. */ +export const isContextLengthError = (e: unknown): boolean => + (errorMessage(e) ?? '').toLowerCase().includes('maximum context length'); + +const firstMatch = ( + message: string, + patterns: RegExp[], +): number | undefined => { + for (const pattern of patterns) { + const value = Number(message.match(pattern)?.[1]); + if (Number.isFinite(value)) return value; + } + return undefined; +}; + +/** + * The largest output cap that fits the window without exceeding `cap`. + * + * `undefined` when neither the rejection nor a request to measure supplies the + * numbers to size one; below 1 when the prompt alone fills the window and no + * retry can fit. + */ +export const contextLengthRetryCap = ({ + error, + cap, + contextWindow, + request, +}: { + error: unknown; + /** + * The credit gate's ceiling for this attempt, if it set one; the retry + * never exceeds it. + */ + cap: number | undefined; + /** The model's declared window, used when the message omits it. */ + contextWindow: number | undefined; + /** The request, measured when the rejection doesn't report its size. */ + request?: { messages?: unknown }; +}): number | undefined => { + const message = errorMessage(error) ?? ''; + + const window = + firstMatch(message, CONTEXT_WINDOW_PATTERNS) ?? contextWindow; + const input = + firstMatch(message, INPUT_TOKEN_PATTERNS) ?? + (request === undefined + ? undefined + : estimatePromptTokens(request.messages ?? []) * + ESTIMATED_PROMPT_MARGIN); + if (window === undefined || input === undefined) return undefined; + + const room = Math.min(cap ?? Number.POSITIVE_INFINITY, window - input); + return Number.isFinite(room) ? Math.floor(room) : undefined; +}; + +/** + * The same request sized to fit, or `undefined` when it can't be — no room + * left, no window to size against, or no cap smaller than the one just rejected + * — and the original rejection stands. Returns a shallow copy; the rejected + * params are still the first attempt's record. + */ +export const contextLengthRetryParams = < + T extends { max_tokens?: number | null; messages?: unknown }, +>( + params: T, + options: { error: unknown; contextWindow: number | undefined }, +): T | undefined => { + const original = params.max_tokens ?? undefined; + const cap = contextLengthRetryCap({ + error: options.error, + cap: original, + contextWindow: options.contextWindow, + request: params, + }); + if (cap === undefined || cap < 1) return undefined; + // The upstream said `original` didn't fit; resending it unchanged only + // earns the same rejection a second time. + if (original !== undefined && cap >= original) return undefined; + return { ...params, max_tokens: cap }; +}; diff --git a/src/backend/drivers/util/tokenEstimate.ts b/src/backend/drivers/util/tokenEstimate.ts index ef6c65810..ad2b5d215 100644 --- a/src/backend/drivers/util/tokenEstimate.ts +++ b/src/backend/drivers/util/tokenEstimate.ts @@ -20,12 +20,13 @@ /** * Tokens a run of text is worth. * - * The v1 estimator: average the two cheap approximations (characters over four, - * words times four thirds) and halve the result. It runs low on purpose — it - * gates and estimates rather than prices, and the real count arrives from the - * provider a moment later. The one estimator is shared by the chat credit gate, - * the chat unreported-stream backstop, and image-prompt pricing, so a future - * calibration moves all of them at once. + * The v1 estimator: the mean of two cheap approximations — characters over + * four, words times four thirds. On prose the two agree and the mean lands + * within a few percent of the real count; on text with few whitespace + * boundaries the word term collapses and the mean runs about half low on JSON + * and code, and a fraction of the real count on base64, logs, or CJK prose. + * Shared by the chat credit gate, the chat unreported-stream backstop, and + * image-prompt pricing, so a calibration moves all of them at once. * * @see https://help.openai.com/en/articles/4936856 */