Added litellm and validator tests

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",
+			)
+		})
+	})
+})

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 {

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()
+	})
+})

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 })
 	}