From 566298d89c520bc1e2203cf7781db2b9b41dcccc Mon Sep 17 00:00:00 2001 From: Neal Shah <30693865+ProgrammerIn-wonderland@users.noreply.github.com> Date: Mon, 10 Aug 2026 13:14:33 -0400 Subject: [PATCH] make context length on togetherAI more lenient (#3533) --- .../TogetherAIProvider.integration.test.ts | 28 +++++ .../together/TogetherAIProvider.test.ts | 114 ++++++++++++++++++ .../providers/together/TogetherAIProvider.ts | 44 ++++++- 3 files changed, 183 insertions(+), 3 deletions(-) diff --git a/src/backend/drivers/ai-chat/providers/together/TogetherAIProvider.integration.test.ts b/src/backend/drivers/ai-chat/providers/together/TogetherAIProvider.integration.test.ts index 0adc36063..93cc3d22b 100644 --- a/src/backend/drivers/ai-chat/providers/together/TogetherAIProvider.integration.test.ts +++ b/src/backend/drivers/ai-chat/providers/together/TogetherAIProvider.integration.test.ts @@ -63,5 +63,33 @@ describe.skipIf(skipUnlessEnv(ENV_VAR))( ?.content; expect(typeof text === 'string' && text.length > 0).toBe(true); }); + + it('recovers when max_tokens leaves no room for the prompt', { timeout: INTEGRATION_TEST_TIMEOUT_MS }, async () => { + const provider = new TogetherAIProvider( + { apiKey: optionalEnv(ENV_VAR)! }, + makeMeteringStub(), + ); + + // Asking for the model's whole context as output leaves no room + // for the prompt, which Together rejects outright — the provider + // should retry uncapped rather than surface a 400. + const model = (await provider.models()).find( + (m) => m.id === 'togetherai:Qwen/Qwen2.5-7B-Instruct-Turbo', + )!; + + const result = await withTestActor(() => + provider.complete({ + model: model.id, + messages: [ + { role: 'user', content: 'Say hi in one word.' }, + ], + max_tokens: model.context!, + }), + ); + + const text = (result as { message?: { content?: string } }).message + ?.content; + expect(typeof text === 'string' && text.length > 0).toBe(true); + }); }, ); 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 9a6c5e94e..bcac330a4 100644 --- a/src/backend/drivers/ai-chat/providers/together/TogetherAIProvider.test.ts +++ b/src/backend/drivers/ai-chat/providers/together/TogetherAIProvider.test.ts @@ -209,6 +209,17 @@ describe('TogetherAIProvider model catalog', () => { expect(ids).toContain('model-fallback-test-1'); }); + it('reserves headroom under the context length for the output cap', async () => { + const { provider } = makeProvider(); + const model = (await provider.models()).find( + (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. + expect(model.context).toBe(32768); + expect(model.max_tokens).toBe(Math.floor(32768 * 0.95)); + }); + it('caches the coerced model list in kv after the first call', async () => { const { provider } = makeProvider(); await provider.models(); @@ -357,6 +368,109 @@ describe('TogetherAIProvider.complete request shape', () => { expect(createMock).not.toHaveBeenCalled(); expect(recordSpy).not.toHaveBeenCalled(); }); + + // Together's APIError carries the whole response body on `.error`, so the + // provider message sits at `.error.error.message` — one level deeper than + // the OpenAI SDK puts it. + const contextLengthError = { + status: 400, + message: + '400 {"id":"ovG6YRd-6z2FuN","error":{"message":"Failed to start generation: The input token count (11) plus the requested output count (1048573) exceeds the model\'s maximum context length (1048576)"}}', + error: { + id: 'ovG6YRd-6z2FuN', + error: { + message: + "Failed to start generation: The input token count (11) plus the requested output count (1048573) exceeds the model's maximum context length (1048576)", + type: 'invalid_request_error', + }, + }, + }; + + it('retries without max_tokens when Together rejects with a context-length error', async () => { + const { provider } = makeProvider(); + + createMock + .mockRejectedValueOnce(contextLengthError) + .mockResolvedValueOnce(baseCompletion); + + await withTestActor(() => + provider.complete({ + model: 'togetherai:Qwen/Qwen2.5-7B-Instruct-Turbo', + messages: [{ role: 'user', content: 'hi' }], + max_tokens: 1048573, + }), + ); + + // 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); + }); + + it('retries a streaming request the same way', async () => { + const { provider } = makeProvider(); + + createMock + .mockRejectedValueOnce(contextLengthError) + .mockReturnValueOnce(asAsyncIterable([])); + + await withTestActor(() => + provider.complete({ + model: 'togetherai:Qwen/Qwen2.5-7B-Instruct-Turbo', + messages: [{ role: 'user', content: 'hi' }], + max_tokens: 1048573, + stream: true, + }), + ); + + expect(createMock).toHaveBeenCalledTimes(2); + expect('max_tokens' in createMock.mock.calls[1]![0]).toBe(false); + expect(createMock.mock.calls[1]![0].stream).toBe(true); + }); + + it('rethrows non-context-length errors without retrying', async () => { + const { provider } = makeProvider(); + const apiError = { + status: 401, + error: { error: { message: 'Invalid API key' } }, + }; + createMock.mockRejectedValueOnce(apiError); + + await expect( + withTestActor(() => + provider.complete({ + model: 'togetherai:Qwen/Qwen2.5-7B-Instruct-Turbo', + messages: [{ role: 'user', content: 'boom' }], + max_tokens: 100, + }), + ), + ).rejects.toBe(apiError); + + expect(createMock).toHaveBeenCalledTimes(1); + expect(recordSpy).not.toHaveBeenCalled(); + }); + + it('rethrows a bodyless connection error instead of failing to read it', async () => { + const { provider } = makeProvider(); + const connErr = new Error('Connection error.'); + createMock.mockRejectedValueOnce(connErr); + + await expect( + withTestActor(() => + provider.complete({ + model: 'togetherai:Qwen/Qwen2.5-7B-Instruct-Turbo', + messages: [{ role: 'user', content: 'hi' }], + max_tokens: 100, + }), + ), + ).rejects.toBe(connErr); + + expect(createMock).toHaveBeenCalledTimes(1); + }); }); // ── Model resolution ──────────────────────────────────────────────── diff --git a/src/backend/drivers/ai-chat/providers/together/TogetherAIProvider.ts b/src/backend/drivers/ai-chat/providers/together/TogetherAIProvider.ts index 4a7561ee7..925029a3d 100644 --- a/src/backend/drivers/ai-chat/providers/together/TogetherAIProvider.ts +++ b/src/backend/drivers/ai-chat/providers/together/TogetherAIProvider.ts @@ -29,6 +29,24 @@ const TOGETHER_AI_CHAT_COST_MAP = { 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; @@ -83,7 +101,13 @@ export class TogetherAIProvider implements IChatProvider { ), ), }, - max_tokens: model.context_length ?? 8000, + // 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. + max_tokens: model.context_length + ? Math.floor(model.context_length * 0.95) + : 8000, }); } } @@ -145,7 +169,7 @@ export class TogetherAIProvider implements IChatProvider { messages = await OpenAIUtil.process_input_messages(messages); - const completion = await this.#together.chat.completions.create({ + const completionParams = { model: modelIdForParams, messages, stream, @@ -153,7 +177,21 @@ export class TogetherAIProvider implements IChatProvider { ...(max_tokens !== undefined ? { max_tokens } : {}), ...(temperature !== undefined ? { temperature } : {}), ...(stream ? { stream_options: { include_usage: true } } : {}), - } as Together.Chat.Completions.CompletionCreateParamsNonStreaming); + } as Together.Chat.Completions.CompletionCreateParamsNonStreaming; + + let completion; + try { + 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. + if (!isContextLengthError(e)) throw e; + delete completionParams.max_tokens; + completion = + await this.#together.chat.completions.create(completionParams); + } return OpenAIUtil.handle_completion_output({ usage_calculator: ({ usage }) => {