Skip to content

Commit bde2c3c

Browse files
fix: use max_output_tokens when available in LiteLLM fetcher (#8455)
Co-authored-by: Roo Code <[email protected]>
1 parent 13d20bb commit bde2c3c

File tree

2 files changed

+88
-1
lines changed

2 files changed

+88
-1
lines changed

src/api/providers/fetchers/__tests__/litellm.spec.ts

Lines changed: 87 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -589,4 +589,91 @@ describe("getLiteLLMModels", () => {
589589

590590
const result = await getLiteLLMModels("test-api-key", "http://localhost:4000")
591591
})
592+
593+
it("prefers max_output_tokens over max_tokens when both are present", async () => {
594+
const mockResponse = {
595+
data: {
596+
data: [
597+
{
598+
model_name: "claude-3-5-sonnet-4-5",
599+
model_info: {
600+
max_tokens: 200000, // This should be ignored
601+
max_output_tokens: 64000, // This should be used
602+
max_input_tokens: 200000,
603+
supports_vision: true,
604+
supports_prompt_caching: false,
605+
supports_computer_use: true,
606+
},
607+
litellm_params: {
608+
model: "anthropic/claude-3-5-sonnet-4-5",
609+
},
610+
},
611+
{
612+
model_name: "model-with-only-max-tokens",
613+
model_info: {
614+
max_tokens: 8192, // This should be used as fallback
615+
// No max_output_tokens
616+
max_input_tokens: 128000,
617+
supports_vision: false,
618+
},
619+
litellm_params: {
620+
model: "test/model-with-only-max-tokens",
621+
},
622+
},
623+
{
624+
model_name: "model-with-only-max-output-tokens",
625+
model_info: {
626+
// No max_tokens
627+
max_output_tokens: 16384, // This should be used
628+
max_input_tokens: 100000,
629+
supports_vision: false,
630+
},
631+
litellm_params: {
632+
model: "test/model-with-only-max-output-tokens",
633+
},
634+
},
635+
],
636+
},
637+
}
638+
639+
mockedAxios.get.mockResolvedValue(mockResponse)
640+
641+
const result = await getLiteLLMModels("test-api-key", "http://localhost:4000")
642+
643+
// Should use max_output_tokens (64000) instead of max_tokens (200000)
644+
expect(result["claude-3-5-sonnet-4-5"]).toEqual({
645+
maxTokens: 64000,
646+
contextWindow: 200000,
647+
supportsImages: true,
648+
supportsComputerUse: true,
649+
supportsPromptCache: false,
650+
inputPrice: undefined,
651+
outputPrice: undefined,
652+
description: "claude-3-5-sonnet-4-5 via LiteLLM proxy",
653+
})
654+
655+
// Should fall back to max_tokens when max_output_tokens is not present
656+
expect(result["model-with-only-max-tokens"]).toEqual({
657+
maxTokens: 8192,
658+
contextWindow: 128000,
659+
supportsImages: false,
660+
supportsComputerUse: false,
661+
supportsPromptCache: false,
662+
inputPrice: undefined,
663+
outputPrice: undefined,
664+
description: "model-with-only-max-tokens via LiteLLM proxy",
665+
})
666+
667+
// Should use max_output_tokens when max_tokens is not present
668+
expect(result["model-with-only-max-output-tokens"]).toEqual({
669+
maxTokens: 16384,
670+
contextWindow: 100000,
671+
supportsImages: false,
672+
supportsComputerUse: false,
673+
supportsPromptCache: false,
674+
inputPrice: undefined,
675+
outputPrice: undefined,
676+
description: "model-with-only-max-output-tokens via LiteLLM proxy",
677+
})
678+
})
592679
})

src/api/providers/fetchers/litellm.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,7 @@ export async function getLiteLLMModels(apiKey: string, baseUrl: string): Promise
4141
if (!modelName || !modelInfo || !litellmModelName) continue
4242

4343
models[modelName] = {
44-
maxTokens: modelInfo.max_tokens || 8192,
44+
maxTokens: modelInfo.max_output_tokens || modelInfo.max_tokens || 8192,
4545
contextWindow: modelInfo.max_input_tokens || 200000,
4646
supportsImages: Boolean(modelInfo.supports_vision),
4747
// litellm_params.model may have a prefix like openrouter/

0 commit comments

Comments
 (0)