diff --git a/packages/types/src/provider-settings.ts b/packages/types/src/provider-settings.ts index 443cbafcf1..22c5cb81ff 100644 --- a/packages/types/src/provider-settings.ts +++ b/packages/types/src/provider-settings.ts @@ -374,6 +374,7 @@ const litellmSchema = baseProviderSettingsSchema.extend({ litellmApiKey: z.string().optional(), litellmModelId: z.string().optional(), litellmUsePromptCache: z.boolean().optional(), + litellmUseAzureBedrock: z.boolean().optional(), }) const cerebrasSchema = apiModelIdProviderModelSchema.extend({ diff --git a/src/api/providers/__tests__/lite-llm.spec.ts b/src/api/providers/__tests__/lite-llm.spec.ts index 80a52d2d4d..7d7c945c83 100644 --- a/src/api/providers/__tests__/lite-llm.spec.ts +++ b/src/api/providers/__tests__/lite-llm.spec.ts @@ -401,7 +401,117 @@ describe("LiteLLMHandler", () => { }) describe("Bedrock model handling", () => { - it("should exclude parallel_tool_calls for Bedrock models when using native tools", async () => { + it("should exclude parallel_tool_calls when litellmUseAzureBedrock is explicitly true", async () => { + const options: ApiHandlerOptions = { + ...mockOptions, + litellmModelId: "gpt-4", // Non-Bedrock model ID + litellmUseAzureBedrock: true, // Explicitly set to Bedrock + } + handler = new LiteLLMHandler(options) + + const systemPrompt = "You are a helpful assistant" + const messages: Anthropic.Messages.MessageParam[] = [{ role: "user", content: "Test" }] + + // Mock the stream response + const mockStream = { + async *[Symbol.asyncIterator]() { + yield { + choices: [{ delta: { content: "Response" } }], + usage: { + prompt_tokens: 10, + completion_tokens: 5, + }, + } + }, + } + + mockCreate.mockReturnValue({ + withResponse: vi.fn().mockResolvedValue({ data: mockStream }), + }) + + const metadata = { + taskId: "test-task", + tools: [ + { + type: "function" as const, + function: { + name: "test_tool", + description: "A test tool", + parameters: { type: "object", properties: {} }, + }, + }, + ], + toolProtocol: "native" as const, + parallelToolCalls: true, + } + + const generator = handler.createMessage(systemPrompt, messages, metadata) + for await (const chunk of generator) { + // Consume the generator + } + + // Verify that parallel_tool_calls is NOT included when litellmUseAzureBedrock is true + const createCall = mockCreate.mock.calls[0][0] + expect(createCall.parallel_tool_calls).toBeUndefined() + expect(createCall.tools).toBeDefined() + }) + + it("should include parallel_tool_calls when litellmUseAzureBedrock is explicitly false", async () => { + const options: ApiHandlerOptions = { + ...mockOptions, + litellmModelId: "bedrock/anthropic.claude-3-5-sonnet-20241022-v2:0", // Bedrock model ID + litellmUseAzureBedrock: false, // Explicitly set to NOT Bedrock + } + handler = new LiteLLMHandler(options) + + const systemPrompt = "You are a helpful assistant" + const messages: Anthropic.Messages.MessageParam[] = [{ role: "user", content: "Test" }] + + // Mock the stream response + const mockStream = { + async *[Symbol.asyncIterator]() { + yield { + choices: [{ delta: { content: "Response" } }], + usage: { + prompt_tokens: 10, + completion_tokens: 5, + }, + } + }, + } + + mockCreate.mockReturnValue({ + withResponse: vi.fn().mockResolvedValue({ data: mockStream }), + }) + + const metadata = { + taskId: "test-task", + tools: [ + { + type: "function" as const, + function: { + name: "test_tool", + description: "A test tool", + parameters: { type: "object", properties: {} }, + }, + }, + ], + toolProtocol: "native" as const, + parallelToolCalls: true, + } + + const generator = handler.createMessage(systemPrompt, messages, metadata) + for await (const chunk of generator) { + // Consume the generator + } + + // Verify that parallel_tool_calls IS included when litellmUseAzureBedrock is false + const createCall = mockCreate.mock.calls[0][0] + expect(createCall.parallel_tool_calls).toBe(true) + expect(createCall.tools).toBeDefined() + }) + + it("should auto-detect and exclude parallel_tool_calls for Bedrock models when litellmUseAzureBedrock is not set", async () => { const bedrockModels = ["bedrock/anthropic.claude-3-5-sonnet-20241022-v2:0", "amazon.titan-text-express-v1"] for (const modelId of bedrockModels) { @@ -461,7 +571,7 @@ describe("LiteLLMHandler", () => { } }) - it("should include parallel_tool_calls for non-Bedrock models when using native tools", async () => { + it("should auto-detect and include parallel_tool_calls for non-Bedrock models when litellmUseAzureBedrock is not set", async () => { const nonBedrockModels = [ "gpt-4", "claude-3-opus", diff --git a/src/api/providers/lite-llm.ts b/src/api/providers/lite-llm.ts index 0ec4418c4a..6bf79c29d6 100644 --- a/src/api/providers/lite-llm.ts +++ b/src/api/providers/lite-llm.ts @@ -41,9 +41,18 @@ export class LiteLLMHandler extends RouterProvider implements SingleCompletionHa /** * Check if the model is routed through AWS Bedrock * Bedrock doesn't support the parallel_tool_calls parameter + * + * If the user has explicitly set litellmUseAzureBedrock, use that setting. + * Otherwise, fall back to auto-detection based on model ID patterns. * Note: We exclude 'anthropic.' prefix as it can match direct Anthropic API access through LiteLLM */ private isBedrockModel(modelId: string): boolean { + // User-specified option takes precedence + if (this.options.litellmUseAzureBedrock !== undefined) { + return this.options.litellmUseAzureBedrock + } + + // Fall back to auto-detection const lowerModel = modelId.toLowerCase() return ( lowerModel.includes("bedrock") || diff --git a/webview-ui/src/components/settings/providers/LiteLLM.tsx b/webview-ui/src/components/settings/providers/LiteLLM.tsx index 0b89b671ce..e08b3c5cb1 100644 --- a/webview-ui/src/components/settings/providers/LiteLLM.tsx +++ b/webview-ui/src/components/settings/providers/LiteLLM.tsx @@ -156,6 +156,20 @@ export const LiteLLM = ({ simplifySettings={simplifySettings} /> + {/* Bedrock backend option */} +
+ { + setApiConfigurationField("litellmUseAzureBedrock", e.target.checked) + }}> + {t("settings:providers.litellmUseAzureBedrock")} + +
+ {t("settings:providers.litellmUseAzureBedrockDescription")} +
+
+ {/* Show prompt caching option if the selected model supports it */} {(() => { const selectedModelId = apiConfiguration.litellmModelId || litellmDefaultModelId diff --git a/webview-ui/src/i18n/locales/en/settings.json b/webview-ui/src/i18n/locales/en/settings.json index 2396bded68..66ae4201ad 100644 --- a/webview-ui/src/i18n/locales/en/settings.json +++ b/webview-ui/src/i18n/locales/en/settings.json @@ -370,6 +370,8 @@ "getXaiApiKey": "Get xAI API Key", "litellmApiKey": "LiteLLM API Key", "litellmBaseUrl": "LiteLLM Base URL", + "litellmUseAzureBedrock": "Backend is AWS Bedrock", + "litellmUseAzureBedrockDescription": "Enable this if your LiteLLM proxy routes to AWS Bedrock models. This ensures Bedrock-incompatible parameters are excluded from requests.", "awsCredentials": "AWS Credentials", "awsProfile": "AWS Profile", "awsApiKey": "Amazon Bedrock API Key",