diff --git a/src/api/providers/__tests__/litellm.test.ts b/src/api/providers/__tests__/litellm.test.ts new file mode 100644 index 0000000000..bc20156d29 --- /dev/null +++ b/src/api/providers/__tests__/litellm.test.ts @@ -0,0 +1,249 @@ +// npx jest src/api/providers/__tests__/litellm.test.ts + +import { Anthropic } from "@anthropic-ai/sdk" // For message types +import OpenAI from "openai" + +import { LiteLLMHandler } from "../litellm" +import { ApiHandlerOptions, litellmDefaultModelId, litellmDefaultModelInfo, ModelInfo } from "../../../shared/api" +import * as modelCache from "../fetchers/modelCache" + +const mockOpenAICreateCompletions = jest.fn() +jest.mock("openai", () => { + return jest.fn(() => ({ + chat: { + completions: { + create: mockOpenAICreateCompletions, + }, + }, + })) +}) + +jest.mock("../fetchers/modelCache", () => ({ + getModels: jest.fn(), +})) + +const mockGetModels = modelCache.getModels as jest.Mock + +describe("LiteLLMHandler", () => { + const defaultMockOptions: ApiHandlerOptions = { + litellmApiKey: "test-litellm-key", + litellmModelId: "litellm-test-model", + litellmBaseUrl: "http://mock-litellm-server:8000", + modelTemperature: 0.1, // Add a default temperature for tests + } + + const mockModelInfo: ModelInfo = { + maxTokens: 4096, + contextWindow: 128000, + supportsImages: false, + supportsPromptCache: true, + supportsComputerUse: false, + description: "A test LiteLLM model", + } + + beforeEach(() => { + jest.clearAllMocks() + + mockGetModels.mockResolvedValue({ + [defaultMockOptions.litellmModelId!]: mockModelInfo, + }) + // Spy on supportsTemperature and default to true for most tests, can be overridden + jest.spyOn(LiteLLMHandler.prototype as any, "supportsTemperature").mockReturnValue(true) + }) + + describe("constructor", () => { + it("initializes with correct options and defaults", () => { + const handler = new LiteLLMHandler(defaultMockOptions) // This will call new OpenAI() + expect(handler).toBeInstanceOf(LiteLLMHandler) + // Check if the mock constructor was called with the right params + expect(OpenAI).toHaveBeenCalledWith({ + baseURL: defaultMockOptions.litellmBaseUrl, + apiKey: defaultMockOptions.litellmApiKey, + }) + }) + + it("uses default baseURL if not provided", () => { + new LiteLLMHandler({ litellmApiKey: "key", litellmModelId: "id" }) + expect(OpenAI).toHaveBeenCalledWith(expect.objectContaining({ baseURL: "http://localhost:4000" })) + }) + + it("uses dummy API key if not provided", () => { + new LiteLLMHandler({ litellmBaseUrl: "url", litellmModelId: "id" }) + expect(OpenAI).toHaveBeenCalledWith(expect.objectContaining({ apiKey: "dummy-key" })) + }) + }) + + describe("fetchModel", () => { + it("returns correct model info when modelId is provided and found in getModels", async () => { + const handler = new LiteLLMHandler(defaultMockOptions) + const result = await handler.fetchModel() + expect(mockGetModels).toHaveBeenCalledWith( + "litellm", + defaultMockOptions.litellmApiKey, + defaultMockOptions.litellmBaseUrl, + ) + expect(result).toEqual({ id: defaultMockOptions.litellmModelId, info: mockModelInfo }) + }) + + it("returns defaultModelInfo if provided modelId is NOT found in getModels result", async () => { + mockGetModels.mockResolvedValueOnce({ "another-model": { contextWindow: 1, supportsPromptCache: false } }) + const handler = new LiteLLMHandler(defaultMockOptions) + const result = await handler.fetchModel() + expect(result.id).toBe(litellmDefaultModelId) + expect(result.info).toEqual(litellmDefaultModelInfo) + }) + + it("uses defaultModelId and its info if litellmModelId option is undefined and defaultModelId is in getModels", async () => { + const specificDefaultModelInfo = { ...mockModelInfo, description: "Specific Default Model Info" } + mockGetModels.mockResolvedValueOnce({ [litellmDefaultModelId]: specificDefaultModelInfo }) + const handler = new LiteLLMHandler({ ...defaultMockOptions, litellmModelId: undefined }) + const result = await handler.fetchModel() + expect(result.id).toBe(litellmDefaultModelId) + expect(result.info).toEqual(specificDefaultModelInfo) + }) + + it("uses defaultModelId and defaultModelInfo if litellmModelId option is undefined and defaultModelId is NOT in getModels", async () => { + mockGetModels.mockResolvedValueOnce({ "some-other-model": mockModelInfo }) + const handler = new LiteLLMHandler({ ...defaultMockOptions, litellmModelId: undefined }) + const result = await handler.fetchModel() + expect(result.id).toBe(litellmDefaultModelId) + expect(result.info).toEqual(litellmDefaultModelInfo) + }) + + it("throws an error if getModels fails", async () => { + mockGetModels.mockRejectedValueOnce(new Error("Network error")) + const handler = new LiteLLMHandler(defaultMockOptions) + await expect(handler.fetchModel()).rejects.toThrow("Network error") + }) + }) + + describe("createMessage", () => { + const systemPrompt = "You are a helpful assistant." + const messages: Anthropic.Messages.MessageParam[] = [{ role: "user", content: "Hello" }] + // mockCreateGlobal is no longer needed here, use mockOpenAICreateCompletions directly + + beforeEach(() => { + // mockOpenAICreateCompletions is already cleared by jest.clearAllMocks() in the outer beforeEach + // or mockOpenAICreateCompletions.mockClear() if we want to be very specific + }) + + it("streams text and usage chunks correctly", async () => { + const mockStreamData = { + async *[Symbol.asyncIterator]() { + yield { id: "chunk1", choices: [{ delta: { content: "Response part 1" } }], usage: null } + yield { id: "chunk2", choices: [{ delta: { content: " part 2" } }], usage: null } + yield { id: "chunk3", choices: [{ delta: {} }], usage: { prompt_tokens: 10, completion_tokens: 5 } } + }, + } + mockOpenAICreateCompletions.mockReturnValue({ + withResponse: jest.fn().mockResolvedValue({ data: mockStreamData }), + }) + + const handler = new LiteLLMHandler(defaultMockOptions) + const generator = handler.createMessage(systemPrompt, messages) + const chunks = [] + for await (const chunk of generator) { + chunks.push(chunk) + } + + expect(chunks).toEqual([ + { type: "text", text: "Response part 1" }, + { type: "text", text: " part 2" }, + { type: "usage", inputTokens: 10, outputTokens: 5 }, + ]) + expect(mockOpenAICreateCompletions).toHaveBeenCalledWith({ + model: defaultMockOptions.litellmModelId, + max_tokens: mockModelInfo.maxTokens, + messages: [ + { role: "system", content: systemPrompt }, + { role: "user", content: "Hello" }, + ], + stream: true, + stream_options: { include_usage: true }, + temperature: defaultMockOptions.modelTemperature, + }) + }) + + it("handles temperature option if supported", async () => { + const handler = new LiteLLMHandler({ ...defaultMockOptions, modelTemperature: 0.7 }) + const mockStreamData = { async *[Symbol.asyncIterator]() {} } + mockOpenAICreateCompletions.mockReturnValue({ + withResponse: jest.fn().mockResolvedValue({ data: mockStreamData }), + }) + + const generator = handler.createMessage(systemPrompt, messages) + for await (const _ of generator) { + } + + expect(mockOpenAICreateCompletions).toHaveBeenCalledWith(expect.objectContaining({ temperature: 0.7 })) + }) + + it("does not include temperature if not supported by model", async () => { + ;(LiteLLMHandler.prototype as any).supportsTemperature.mockReturnValue(false) + const handler = new LiteLLMHandler(defaultMockOptions) + const mockStreamData = { async *[Symbol.asyncIterator]() {} } + mockOpenAICreateCompletions.mockReturnValue({ + withResponse: jest.fn().mockResolvedValue({ data: mockStreamData }), + }) + + const generator = handler.createMessage(systemPrompt, messages) + for await (const _ of generator) { + } + + const callArgs = mockOpenAICreateCompletions.mock.calls[0][0] + expect(callArgs.temperature).toBeUndefined() + }) + + it("throws a formatted error if API call (streaming) fails", async () => { + const apiError = new Error("LLM Provider Error") + // Simulate the error occurring within the stream itself + mockOpenAICreateCompletions.mockReturnValue({ + withResponse: jest.fn().mockResolvedValue({ + data: { + async *[Symbol.asyncIterator]() { + throw apiError + }, + }, + }), + }) + + const handler = new LiteLLMHandler(defaultMockOptions) + const generator = handler.createMessage(systemPrompt, messages) + await expect(async () => { + for await (const _ of generator) { + } + }).rejects.toThrow("LiteLLM streaming error: " + apiError.message) + }) + }) + + describe("completePrompt", () => { + const prompt = "Translate 'hello' to French." + // mockCreateGlobal is no longer needed here, use mockOpenAICreateCompletions directly + + beforeEach(() => { + // mockOpenAICreateCompletions is already cleared by jest.clearAllMocks() in the outer beforeEach + }) + + it("returns completion successfully", async () => { + mockOpenAICreateCompletions.mockResolvedValueOnce({ choices: [{ message: { content: "Bonjour" } }] }) + const handler = new LiteLLMHandler(defaultMockOptions) + const result = await handler.completePrompt(prompt) + + expect(result).toBe("Bonjour") + expect(mockOpenAICreateCompletions).toHaveBeenCalledWith({ + model: defaultMockOptions.litellmModelId, + max_tokens: mockModelInfo.maxTokens, + messages: [{ role: "user", content: prompt }], + temperature: defaultMockOptions.modelTemperature, + }) + }) + + it("throws a formatted error if API call fails", async () => { + mockOpenAICreateCompletions.mockRejectedValueOnce(new Error("Completion API Down")) + const handler = new LiteLLMHandler(defaultMockOptions) + await expect(handler.completePrompt(prompt)).rejects.toThrow( + "LiteLLM completion error: Completion API Down", + ) + }) + }) +}) diff --git a/src/api/providers/router-provider.ts b/src/api/providers/router-provider.ts index a0decdcab4..37557a3af0 100644 --- a/src/api/providers/router-provider.ts +++ b/src/api/providers/router-provider.ts @@ -17,6 +17,8 @@ type RouterProviderOptions = { export abstract class RouterProvider extends BaseProvider { protected readonly options: ApiHandlerOptions protected readonly name: RouterName + protected readonly baseURL: string + protected readonly apiKey: string protected models: ModelRecord = {} protected readonly modelId?: string protected readonly defaultModelId: string @@ -39,21 +41,37 @@ export abstract class RouterProvider extends BaseProvider { this.modelId = modelId this.defaultModelId = defaultModelId this.defaultModelInfo = defaultModelInfo + this.baseURL = baseURL + this.apiKey = apiKey - this.client = new OpenAI({ baseURL, apiKey }) + this.client = new OpenAI({ + baseURL, + apiKey, + }) } public async fetchModel() { - this.models = await getModels(this.name, this.client.apiKey, this.client.baseURL) + this.models = await getModels(this.name, this.apiKey, this.baseURL) return this.getModel() } override getModel(): { id: string; info: ModelInfo } { - const id = this.modelId ?? this.defaultModelId + const userSpecifiedModelId = this.modelId - return this.models[id] - ? { id, info: this.models[id] } - : { id: this.defaultModelId, info: this.defaultModelInfo } + // Priority 1: Use user-specified model if it's valid and found in fetched models + if (userSpecifiedModelId && this.models[userSpecifiedModelId]) { + return { id: userSpecifiedModelId, info: this.models[userSpecifiedModelId] } + } + + // Priority 2: If user-specified model is not found (or not specified at all), + // try the provider's default model ID with its fetched info (if available). + if (this.models[this.defaultModelId]) { + return { id: this.defaultModelId, info: this.models[this.defaultModelId] } + } + + // Priority 3: Ultimate fallback: provider's default model ID with its (static) defaultModelInfo. + // This is reached if userSpecifiedModelId was invalid/not found AND this.defaultModelId was also not in this.models. + return { id: this.defaultModelId, info: this.defaultModelInfo } } protected supportsTemperature(modelId: string): boolean { diff --git a/webview-ui/src/utils/__tests__/validate.test.ts b/webview-ui/src/utils/__tests__/validate.test.ts new file mode 100644 index 0000000000..8faa0743a7 --- /dev/null +++ b/webview-ui/src/utils/__tests__/validate.test.ts @@ -0,0 +1,73 @@ +// npx jest webview-ui/src/utils/__tests__/validate.test.ts + +import { validateModelId } from "../validate" +import { ProviderSettings, RouterModels } from "@roo/shared/api" + +// Mock i18next.t for error messages +jest.mock("i18next", () => ({ + t: (key: string, opts?: any) => { + if (key === "settings:validation.modelAvailability") { + return `Model ${opts.modelId} not available` + } + if (key === "settings:validation.modelId") { + return "Model ID required" + } + return key + }, +})) + +describe("validateModelId", () => { + const baseConfig: ProviderSettings = { + apiProvider: "litellm", + litellmModelId: "foo-model", + litellmApiKey: "key", + litellmBaseUrl: "http://localhost:4000", + } as any + + it("returns undefined if model is in the list", () => { + const routerModels: RouterModels = { + litellm: { "foo-model": { contextWindow: 1, supportsPromptCache: false } }, + openrouter: {}, + glama: {}, + unbound: {}, + requesty: {}, + } + expect(validateModelId(baseConfig, routerModels)).toBeUndefined() + }) + + it("returns error if model is not in the list", () => { + const routerModels: RouterModels = { + litellm: { "another-model": { contextWindow: 1, supportsPromptCache: false } }, + openrouter: {}, + glama: {}, + unbound: {}, + requesty: {}, + } + expect(validateModelId(baseConfig, routerModels)).toBe("Model foo-model not available") + }) + + it("returns undefined if routerModels is undefined", () => { + expect(validateModelId(baseConfig, undefined)).toBeUndefined() + }) + + it("returns error if modelId is missing", () => { + const config = { ...baseConfig, litellmModelId: undefined } + expect(validateModelId(config, undefined)).toBe("Model ID required") + }) + + it("returns error if model list is empty", () => { + const routerModels: RouterModels = { + litellm: {}, + openrouter: {}, + glama: {}, + unbound: {}, + requesty: {}, + } + expect(validateModelId(baseConfig, routerModels)).toBe("Model foo-model not available") + }) + + it("returns undefined for non-router providers", () => { + const config: ProviderSettings = { ...baseConfig, apiProvider: "openai" } + expect(validateModelId(config, undefined)).toBeUndefined() + }) +}) diff --git a/webview-ui/src/utils/validate.ts b/webview-ui/src/utils/validate.ts index 0765fffeda..68f755485d 100644 --- a/webview-ui/src/utils/validate.ts +++ b/webview-ui/src/utils/validate.ts @@ -151,7 +151,7 @@ export function validateModelId(apiConfiguration: ProviderSettings, routerModels const models = routerModels?.[provider] - if (models && Object.keys(models).length > 1 && !Object.keys(models).includes(modelId)) { + if (models && !Object.keys(models).includes(modelId)) { return i18next.t("settings:validation.modelAvailability", { modelId }) }