diff --git a/src/api/providers/__tests__/chutes.spec.ts b/src/api/providers/__tests__/chutes.spec.ts index 911c848b12..cc018a6f87 100644 --- a/src/api/providers/__tests__/chutes.spec.ts +++ b/src/api/providers/__tests__/chutes.spec.ts @@ -325,10 +325,51 @@ describe("ChutesHandler", () => { ) }) - it("createMessage should pass correct parameters to Chutes client for non-DeepSeek models", async () => { + it("createMessage should not include max_tokens by default for non-DeepSeek models", async () => { + const modelId: ChutesModelId = "unsloth/Llama-3.3-70B-Instruct" + const handlerWithModel = new ChutesHandler({ apiModelId: modelId, chutesApiKey: "test-chutes-api-key" }) + + mockCreate.mockImplementationOnce(() => { + return { + [Symbol.asyncIterator]: () => ({ + async next() { + return { done: true } + }, + }), + } + }) + + const systemPrompt = "Test system prompt for Chutes" + const messages: Anthropic.Messages.MessageParam[] = [{ role: "user", content: "Test message for Chutes" }] + + const messageGenerator = handlerWithModel.createMessage(systemPrompt, messages) + await messageGenerator.next() + + expect(mockCreate).toHaveBeenCalledWith( + expect.objectContaining({ + model: modelId, + temperature: 0.5, + messages: expect.arrayContaining([{ role: "system", content: systemPrompt }]), + stream: true, + stream_options: { include_usage: true }, + }), + ) + // Verify max_tokens is NOT included + expect(mockCreate).toHaveBeenCalledWith( + expect.not.objectContaining({ + max_tokens: expect.anything(), + }), + ) + }) + + it("createMessage should include max_tokens when includeMaxTokens is true for non-DeepSeek models", async () => { const modelId: ChutesModelId = "unsloth/Llama-3.3-70B-Instruct" const modelInfo = chutesModels[modelId] - const handlerWithModel = new ChutesHandler({ apiModelId: modelId, chutesApiKey: "test-chutes-api-key" }) + const handlerWithModel = new ChutesHandler({ + apiModelId: modelId, + chutesApiKey: "test-chutes-api-key", + includeMaxTokens: true, + }) mockCreate.mockImplementationOnce(() => { return { diff --git a/src/api/providers/__tests__/fireworks.spec.ts b/src/api/providers/__tests__/fireworks.spec.ts index cfab672c08..5dd227490f 100644 --- a/src/api/providers/__tests__/fireworks.spec.ts +++ b/src/api/providers/__tests__/fireworks.spec.ts @@ -324,12 +324,53 @@ describe("FireworksHandler", () => { expect(firstChunk.value).toEqual({ type: "usage", inputTokens: 10, outputTokens: 20 }) }) - it("createMessage should pass correct parameters to Fireworks client", async () => { + it("createMessage should not include max_tokens by default", async () => { + const modelId: FireworksModelId = "accounts/fireworks/models/kimi-k2-instruct" + const handlerWithModel = new FireworksHandler({ + apiModelId: modelId, + fireworksApiKey: "test-fireworks-api-key", + }) + + mockCreate.mockImplementationOnce(() => { + return { + [Symbol.asyncIterator]: () => ({ + async next() { + return { done: true } + }, + }), + } + }) + + const systemPrompt = "Test system prompt for Fireworks" + const messages: Anthropic.Messages.MessageParam[] = [{ role: "user", content: "Test message for Fireworks" }] + + const messageGenerator = handlerWithModel.createMessage(systemPrompt, messages) + await messageGenerator.next() + + expect(mockCreate).toHaveBeenCalledWith( + expect.objectContaining({ + model: modelId, + temperature: 0.5, + messages: expect.arrayContaining([{ role: "system", content: systemPrompt }]), + stream: true, + stream_options: { include_usage: true }, + }), + ) + // Verify max_tokens is NOT included + expect(mockCreate).toHaveBeenCalledWith( + expect.not.objectContaining({ + max_tokens: expect.anything(), + }), + ) + }) + + it("createMessage should include max_tokens when includeMaxTokens is true", async () => { const modelId: FireworksModelId = "accounts/fireworks/models/kimi-k2-instruct" const modelInfo = fireworksModels[modelId] const handlerWithModel = new FireworksHandler({ apiModelId: modelId, fireworksApiKey: "test-fireworks-api-key", + includeMaxTokens: true, }) mockCreate.mockImplementationOnce(() => { diff --git a/src/api/providers/__tests__/groq.spec.ts b/src/api/providers/__tests__/groq.spec.ts index 72a834b21d..601d01ffcf 100644 --- a/src/api/providers/__tests__/groq.spec.ts +++ b/src/api/providers/__tests__/groq.spec.ts @@ -111,10 +111,51 @@ describe("GroqHandler", () => { expect(firstChunk.value).toEqual({ type: "usage", inputTokens: 10, outputTokens: 20 }) }) - it("createMessage should pass correct parameters to Groq client", async () => { + it("createMessage should not include max_tokens by default", async () => { + const modelId: GroqModelId = "llama-3.1-8b-instant" + const handlerWithModel = new GroqHandler({ apiModelId: modelId, groqApiKey: "test-groq-api-key" }) + + mockCreate.mockImplementationOnce(() => { + return { + [Symbol.asyncIterator]: () => ({ + async next() { + return { done: true } + }, + }), + } + }) + + const systemPrompt = "Test system prompt for Groq" + const messages: Anthropic.Messages.MessageParam[] = [{ role: "user", content: "Test message for Groq" }] + + const messageGenerator = handlerWithModel.createMessage(systemPrompt, messages) + await messageGenerator.next() + + expect(mockCreate).toHaveBeenCalledWith( + expect.objectContaining({ + model: modelId, + temperature: 0.5, + messages: expect.arrayContaining([{ role: "system", content: systemPrompt }]), + stream: true, + stream_options: { include_usage: true }, + }), + ) + // Verify max_tokens is NOT included + expect(mockCreate).toHaveBeenCalledWith( + expect.not.objectContaining({ + max_tokens: expect.anything(), + }), + ) + }) + + it("createMessage should include max_tokens when includeMaxTokens is true", async () => { const modelId: GroqModelId = "llama-3.1-8b-instant" const modelInfo = groqModels[modelId] - const handlerWithModel = new GroqHandler({ apiModelId: modelId, groqApiKey: "test-groq-api-key" }) + const handlerWithModel = new GroqHandler({ + apiModelId: modelId, + groqApiKey: "test-groq-api-key", + includeMaxTokens: true, + }) mockCreate.mockImplementationOnce(() => { return { @@ -143,4 +184,42 @@ describe("GroqHandler", () => { }), ) }) + + it("createMessage should use modelMaxTokens over default when includeMaxTokens is true", async () => { + const modelId: GroqModelId = "llama-3.1-8b-instant" + const customMaxTokens = 2048 + const handlerWithModel = new GroqHandler({ + apiModelId: modelId, + groqApiKey: "test-groq-api-key", + includeMaxTokens: true, + modelMaxTokens: customMaxTokens, + }) + + mockCreate.mockImplementationOnce(() => { + return { + [Symbol.asyncIterator]: () => ({ + async next() { + return { done: true } + }, + }), + } + }) + + const systemPrompt = "Test system prompt for Groq" + const messages: Anthropic.Messages.MessageParam[] = [{ role: "user", content: "Test message for Groq" }] + + const messageGenerator = handlerWithModel.createMessage(systemPrompt, messages) + await messageGenerator.next() + + expect(mockCreate).toHaveBeenCalledWith( + expect.objectContaining({ + model: modelId, + max_tokens: customMaxTokens, + temperature: 0.5, + messages: expect.arrayContaining([{ role: "system", content: systemPrompt }]), + stream: true, + stream_options: { include_usage: true }, + }), + ) + }) }) diff --git a/src/api/providers/__tests__/sambanova.spec.ts b/src/api/providers/__tests__/sambanova.spec.ts index cd0e4a1989..d8d822d449 100644 --- a/src/api/providers/__tests__/sambanova.spec.ts +++ b/src/api/providers/__tests__/sambanova.spec.ts @@ -116,12 +116,53 @@ describe("SambaNovaHandler", () => { expect(firstChunk.value).toEqual({ type: "usage", inputTokens: 10, outputTokens: 20 }) }) - it("createMessage should pass correct parameters to SambaNova client", async () => { + it("createMessage should not include max_tokens by default", async () => { + const modelId: SambaNovaModelId = "Meta-Llama-3.3-70B-Instruct" + const handlerWithModel = new SambaNovaHandler({ + apiModelId: modelId, + sambaNovaApiKey: "test-sambanova-api-key", + }) + + mockCreate.mockImplementationOnce(() => { + return { + [Symbol.asyncIterator]: () => ({ + async next() { + return { done: true } + }, + }), + } + }) + + const systemPrompt = "Test system prompt for SambaNova" + const messages: Anthropic.Messages.MessageParam[] = [{ role: "user", content: "Test message for SambaNova" }] + + const messageGenerator = handlerWithModel.createMessage(systemPrompt, messages) + await messageGenerator.next() + + expect(mockCreate).toHaveBeenCalledWith( + expect.objectContaining({ + model: modelId, + temperature: 0.7, + messages: expect.arrayContaining([{ role: "system", content: systemPrompt }]), + stream: true, + stream_options: { include_usage: true }, + }), + ) + // Verify max_tokens is NOT included + expect(mockCreate).toHaveBeenCalledWith( + expect.not.objectContaining({ + max_tokens: expect.anything(), + }), + ) + }) + + it("createMessage should include max_tokens when includeMaxTokens is true", async () => { const modelId: SambaNovaModelId = "Meta-Llama-3.3-70B-Instruct" const modelInfo = sambaNovaModels[modelId] const handlerWithModel = new SambaNovaHandler({ apiModelId: modelId, sambaNovaApiKey: "test-sambanova-api-key", + includeMaxTokens: true, }) mockCreate.mockImplementationOnce(() => { diff --git a/src/api/providers/__tests__/zai.spec.ts b/src/api/providers/__tests__/zai.spec.ts index 6b93aaa43b..b1e181a230 100644 --- a/src/api/providers/__tests__/zai.spec.ts +++ b/src/api/providers/__tests__/zai.spec.ts @@ -191,13 +191,55 @@ describe("ZAiHandler", () => { expect(firstChunk.value).toEqual({ type: "usage", inputTokens: 10, outputTokens: 20 }) }) - it("createMessage should pass correct parameters to Z AI client", async () => { + it("createMessage should not include max_tokens by default", async () => { + const modelId: InternationalZAiModelId = "glm-4.5" + const handlerWithModel = new ZAiHandler({ + apiModelId: modelId, + zaiApiKey: "test-zai-api-key", + zaiApiLine: "international", + }) + + mockCreate.mockImplementationOnce(() => { + return { + [Symbol.asyncIterator]: () => ({ + async next() { + return { done: true } + }, + }), + } + }) + + const systemPrompt = "Test system prompt for Z AI" + const messages: Anthropic.Messages.MessageParam[] = [{ role: "user", content: "Test message for Z AI" }] + + const messageGenerator = handlerWithModel.createMessage(systemPrompt, messages) + await messageGenerator.next() + + expect(mockCreate).toHaveBeenCalledWith( + expect.objectContaining({ + model: modelId, + temperature: ZAI_DEFAULT_TEMPERATURE, + messages: expect.arrayContaining([{ role: "system", content: systemPrompt }]), + stream: true, + stream_options: { include_usage: true }, + }), + ) + // Verify max_tokens is NOT included + expect(mockCreate).toHaveBeenCalledWith( + expect.not.objectContaining({ + max_tokens: expect.anything(), + }), + ) + }) + + it("createMessage should include max_tokens when includeMaxTokens is true", async () => { const modelId: InternationalZAiModelId = "glm-4.5" const modelInfo = internationalZAiModels[modelId] const handlerWithModel = new ZAiHandler({ apiModelId: modelId, zaiApiKey: "test-zai-api-key", zaiApiLine: "international", + includeMaxTokens: true, }) mockCreate.mockImplementationOnce(() => { diff --git a/src/api/providers/base-openai-compatible-provider.ts b/src/api/providers/base-openai-compatible-provider.ts index f196b5f309..70ccf94401 100644 --- a/src/api/providers/base-openai-compatible-provider.ts +++ b/src/api/providers/base-openai-compatible-provider.ts @@ -67,22 +67,24 @@ export abstract class BaseOpenAiCompatibleProvider messages: Anthropic.Messages.MessageParam[], metadata?: ApiHandlerCreateMessageMetadata, ): ApiStream { - const { - id: model, - info: { maxTokens: max_tokens }, - } = this.getModel() + const { id: model, info } = this.getModel() const temperature = this.options.modelTemperature ?? this.defaultTemperature const params: OpenAI.Chat.Completions.ChatCompletionCreateParamsStreaming = { model, - max_tokens, temperature, messages: [{ role: "system", content: systemPrompt }, ...convertToOpenAiMessages(messages)], stream: true, stream_options: { include_usage: true }, } + // Only add max_tokens if includeMaxTokens is true + if (this.options.includeMaxTokens === true) { + // Use user-configured modelMaxTokens if available, otherwise fall back to model's default maxTokens + params.max_tokens = this.options.modelMaxTokens || info.maxTokens + } + const stream = await this.client.chat.completions.create(params) for await (const chunk of stream) {