Skip to content

Commit 1791bb9

Browse files
slytechnicalslytechnical
authored
Ability to refresh LiteLLM models list (take 2) (#3852)
* Litellm models can now be refreshed * Fix no-case-declarations lint issue and put back missing autoCondenseContextPercent webviewMessageHandler case * Add tests for litellm config changes * replace hardcoded keys with constants
1 parent 8729027 commit 1791bb9

34 files changed

+1331
-111
lines changed

‎src/api/providers/fetchers/__tests__/litellm.test.ts‎

Lines changed: 227 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,227 @@
1+
import axios from "axios"
2+
import { getLiteLLMModels } from "../litellm"
3+
import { COMPUTER_USE_MODELS } from "../../../../shared/api"
4+
5+
// Mock axios
6+
jest.mock("axios")
7+
const mockedAxios = axios as jest.Mocked<typeof axios>
8+
9+
const DUMMY_INVALID_KEY = "invalid-key-for-testing"
10+
11+
describe("getLiteLLMModels", () => {
12+
beforeEach(() => {
13+
jest.clearAllMocks()
14+
})
15+
16+
it("successfully fetches and formats LiteLLM models", async () => {
17+
const mockResponse = {
18+
data: {
19+
data: [
20+
{
21+
model_name: "claude-3-5-sonnet",
22+
model_info: {
23+
max_tokens: 4096,
24+
max_input_tokens: 200000,
25+
supports_vision: true,
26+
supports_prompt_caching: false,
27+
input_cost_per_token: 0.000003,
28+
output_cost_per_token: 0.000015,
29+
},
30+
litellm_params: {
31+
model: "anthropic/claude-3.5-sonnet",
32+
},
33+
},
34+
{
35+
model_name: "gpt-4-turbo",
36+
model_info: {
37+
max_tokens: 8192,
38+
max_input_tokens: 128000,
39+
supports_vision: false,
40+
supports_prompt_caching: false,
41+
input_cost_per_token: 0.00001,
42+
output_cost_per_token: 0.00003,
43+
},
44+
litellm_params: {
45+
model: "openai/gpt-4-turbo",
46+
},
47+
},
48+
],
49+
},
50+
}
51+
52+
mockedAxios.get.mockResolvedValue(mockResponse)
53+
54+
const result = await getLiteLLMModels("test-api-key", "http://localhost:4000")
55+
56+
expect(mockedAxios.get).toHaveBeenCalledWith("http://localhost:4000/v1/model/info", {
57+
headers: {
58+
Authorization: "Bearer test-api-key",
59+
"Content-Type": "application/json",
60+
},
61+
timeout: 5000,
62+
})
63+
64+
expect(result).toEqual({
65+
"claude-3-5-sonnet": {
66+
maxTokens: 4096,
67+
contextWindow: 200000,
68+
supportsImages: true,
69+
supportsComputerUse: true,
70+
supportsPromptCache: false,
71+
inputPrice: 3,
72+
outputPrice: 15,
73+
description: "claude-3-5-sonnet via LiteLLM proxy",
74+
},
75+
"gpt-4-turbo": {
76+
maxTokens: 8192,
77+
contextWindow: 128000,
78+
supportsImages: false,
79+
supportsComputerUse: false,
80+
supportsPromptCache: false,
81+
inputPrice: 10,
82+
outputPrice: 30,
83+
description: "gpt-4-turbo via LiteLLM proxy",
84+
},
85+
})
86+
})
87+
88+
it("makes request without authorization header when no API key provided", async () => {
89+
const mockResponse = {
90+
data: {
91+
data: [],
92+
},
93+
}
94+
95+
mockedAxios.get.mockResolvedValue(mockResponse)
96+
97+
await getLiteLLMModels("", "http://localhost:4000")
98+
99+
expect(mockedAxios.get).toHaveBeenCalledWith("http://localhost:4000/v1/model/info", {
100+
headers: {
101+
"Content-Type": "application/json",
102+
},
103+
timeout: 5000,
104+
})
105+
})
106+
107+
it("handles computer use models correctly", async () => {
108+
const computerUseModel = Array.from(COMPUTER_USE_MODELS)[0]
109+
const mockResponse = {
110+
data: {
111+
data: [
112+
{
113+
model_name: "test-computer-model",
114+
model_info: {
115+
max_tokens: 4096,
116+
max_input_tokens: 200000,
117+
supports_vision: true,
118+
},
119+
litellm_params: {
120+
model: `anthropic/${computerUseModel}`,
121+
},
122+
},
123+
],
124+
},
125+
}
126+
127+
mockedAxios.get.mockResolvedValue(mockResponse)
128+
129+
const result = await getLiteLLMModels("test-api-key", "http://localhost:4000")
130+
131+
expect(result["test-computer-model"]).toEqual({
132+
maxTokens: 4096,
133+
contextWindow: 200000,
134+
supportsImages: true,
135+
supportsComputerUse: true,
136+
supportsPromptCache: false,
137+
inputPrice: undefined,
138+
outputPrice: undefined,
139+
description: "test-computer-model via LiteLLM proxy",
140+
})
141+
})
142+
143+
it("throws error for unexpected response format", async () => {
144+
const mockResponse = {
145+
data: {
146+
// Missing 'data' field
147+
models: [],
148+
},
149+
}
150+
151+
mockedAxios.get.mockResolvedValue(mockResponse)
152+
153+
await expect(getLiteLLMModels("test-api-key", "http://localhost:4000")).rejects.toThrow(
154+
"Failed to fetch LiteLLM models: Unexpected response format.",
155+
)
156+
})
157+
158+
it("throws detailed error for HTTP error responses", async () => {
159+
const axiosError = {
160+
response: {
161+
status: 401,
162+
statusText: "Unauthorized",
163+
},
164+
isAxiosError: true,
165+
}
166+
167+
mockedAxios.isAxiosError.mockReturnValue(true)
168+
mockedAxios.get.mockRejectedValue(axiosError)
169+
170+
await expect(getLiteLLMModels(DUMMY_INVALID_KEY, "http://localhost:4000")).rejects.toThrow(
171+
"Failed to fetch LiteLLM models: 401 Unauthorized. Check base URL and API key.",
172+
)
173+
})
174+
175+
it("throws network error for request failures", async () => {
176+
const axiosError = {
177+
request: {},
178+
isAxiosError: true,
179+
}
180+
181+
mockedAxios.isAxiosError.mockReturnValue(true)
182+
mockedAxios.get.mockRejectedValue(axiosError)
183+
184+
await expect(getLiteLLMModels("test-api-key", "http://invalid-url")).rejects.toThrow(
185+
"Failed to fetch LiteLLM models: No response from server. Check LiteLLM server status and base URL.",
186+
)
187+
})
188+
189+
it("throws generic error for other failures", async () => {
190+
const genericError = new Error("Network timeout")
191+
192+
mockedAxios.isAxiosError.mockReturnValue(false)
193+
mockedAxios.get.mockRejectedValue(genericError)
194+
195+
await expect(getLiteLLMModels("test-api-key", "http://localhost:4000")).rejects.toThrow(
196+
"Failed to fetch LiteLLM models: Network timeout",
197+
)
198+
})
199+
200+
it("handles timeout parameter correctly", async () => {
201+
const mockResponse = { data: { data: [] } }
202+
mockedAxios.get.mockResolvedValue(mockResponse)
203+
204+
await getLiteLLMModels("test-api-key", "http://localhost:4000")
205+
206+
expect(mockedAxios.get).toHaveBeenCalledWith(
207+
"http://localhost:4000/v1/model/info",
208+
expect.objectContaining({
209+
timeout: 5000,
210+
}),
211+
)
212+
})
213+
214+
it("returns empty object when data array is empty", async () => {
215+
const mockResponse = {
216+
data: {
217+
data: [],
218+
},
219+
}
220+
221+
mockedAxios.get.mockResolvedValue(mockResponse)
222+
223+
const result = await getLiteLLMModels("test-api-key", "http://localhost:4000")
224+
225+
expect(result).toEqual({})
226+
})
227+
})

‎src/api/providers/fetchers/__tests__/modelCache.test.ts‎

Lines changed: 158 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,158 @@
1+
import { getModels } from "../modelCache"
2+
import { getLiteLLMModels } from "../litellm"
3+
import { getOpenRouterModels } from "../openrouter"
4+
import { getRequestyModels } from "../requesty"
5+
import { getGlamaModels } from "../glama"
6+
import { getUnboundModels } from "../unbound"
7+
8+
// Mock NodeCache to avoid cache interference
9+
jest.mock("node-cache", () => {
10+
return jest.fn().mockImplementation(() => ({
11+
get: jest.fn().mockReturnValue(undefined), // Always return cache miss
12+
set: jest.fn(),
13+
del: jest.fn(),
14+
}))
15+
})
16+
17+
// Mock fs/promises to avoid file system operations
18+
jest.mock("fs/promises", () => ({
19+
writeFile: jest.fn().mockResolvedValue(undefined),
20+
readFile: jest.fn().mockResolvedValue("{}"),
21+
mkdir: jest.fn().mockResolvedValue(undefined),
22+
}))
23+
24+
// Mock all the model fetchers
25+
jest.mock("../litellm")
26+
jest.mock("../openrouter")
27+
jest.mock("../requesty")
28+
jest.mock("../glama")
29+
jest.mock("../unbound")
30+
31+
const mockGetLiteLLMModels = getLiteLLMModels as jest.MockedFunction<typeof getLiteLLMModels>
32+
const mockGetOpenRouterModels = getOpenRouterModels as jest.MockedFunction<typeof getOpenRouterModels>
33+
const mockGetRequestyModels = getRequestyModels as jest.MockedFunction<typeof getRequestyModels>
34+
const mockGetGlamaModels = getGlamaModels as jest.MockedFunction<typeof getGlamaModels>
35+
const mockGetUnboundModels = getUnboundModels as jest.MockedFunction<typeof getUnboundModels>
36+
37+
const DUMMY_REQUESTY_KEY = "requesty-key-for-testing"
38+
const DUMMY_UNBOUND_KEY = "unbound-key-for-testing"
39+
40+
describe("getModels with new GetModelsOptions", () => {
41+
beforeEach(() => {
42+
jest.clearAllMocks()
43+
})
44+
45+
it("calls getLiteLLMModels with correct parameters", async () => {
46+
const mockModels = {
47+
"claude-3-sonnet": {
48+
maxTokens: 4096,
49+
contextWindow: 200000,
50+
supportsPromptCache: false,
51+
description: "Claude 3 Sonnet via LiteLLM",
52+
},
53+
}
54+
mockGetLiteLLMModels.mockResolvedValue(mockModels)
55+
56+
const result = await getModels({
57+
provider: "litellm",
58+
apiKey: "test-api-key",
59+
baseUrl: "http://localhost:4000",
60+
})
61+
62+
expect(mockGetLiteLLMModels).toHaveBeenCalledWith("test-api-key", "http://localhost:4000")
63+
expect(result).toEqual(mockModels)
64+
})
65+
66+
it("calls getOpenRouterModels for openrouter provider", async () => {
67+
const mockModels = {
68+
"openrouter/model": {
69+
maxTokens: 8192,
70+
contextWindow: 128000,
71+
supportsPromptCache: false,
72+
description: "OpenRouter model",
73+
},
74+
}
75+
mockGetOpenRouterModels.mockResolvedValue(mockModels)
76+
77+
const result = await getModels({ provider: "openrouter" })
78+
79+
expect(mockGetOpenRouterModels).toHaveBeenCalled()
80+
expect(result).toEqual(mockModels)
81+
})
82+
83+
it("calls getRequestyModels with optional API key", async () => {
84+
const mockModels = {
85+
"requesty/model": {
86+
maxTokens: 4096,
87+
contextWindow: 8192,
88+
supportsPromptCache: false,
89+
description: "Requesty model",
90+
},
91+
}
92+
mockGetRequestyModels.mockResolvedValue(mockModels)
93+
94+
const result = await getModels({ provider: "requesty", apiKey: DUMMY_REQUESTY_KEY })
95+
96+
expect(mockGetRequestyModels).toHaveBeenCalledWith(DUMMY_REQUESTY_KEY)
97+
expect(result).toEqual(mockModels)
98+
})
99+
100+
it("calls getGlamaModels for glama provider", async () => {
101+
const mockModels = {
102+
"glama/model": {
103+
maxTokens: 4096,
104+
contextWindow: 8192,
105+
supportsPromptCache: false,
106+
description: "Glama model",
107+
},
108+
}
109+
mockGetGlamaModels.mockResolvedValue(mockModels)
110+
111+
const result = await getModels({ provider: "glama" })
112+
113+
expect(mockGetGlamaModels).toHaveBeenCalled()
114+
expect(result).toEqual(mockModels)
115+
})
116+
117+
it("calls getUnboundModels with optional API key", async () => {
118+
const mockModels = {
119+
"unbound/model": {
120+
maxTokens: 4096,
121+
contextWindow: 8192,
122+
supportsPromptCache: false,
123+
description: "Unbound model",
124+
},
125+
}
126+
mockGetUnboundModels.mockResolvedValue(mockModels)
127+
128+
const result = await getModels({ provider: "unbound", apiKey: DUMMY_UNBOUND_KEY })
129+
130+
expect(mockGetUnboundModels).toHaveBeenCalledWith(DUMMY_UNBOUND_KEY)
131+
expect(result).toEqual(mockModels)
132+
})
133+
134+
it("handles errors and re-throws them", async () => {
135+
const expectedError = new Error("LiteLLM connection failed")
136+
mockGetLiteLLMModels.mockRejectedValue(expectedError)
137+
138+
await expect(
139+
getModels({
140+
provider: "litellm",
141+
apiKey: "test-api-key",
142+
baseUrl: "http://localhost:4000",
143+
}),
144+
).rejects.toThrow("LiteLLM connection failed")
145+
})
146+
147+
it("validates exhaustive provider checking with unknown provider", async () => {
148+
// This test ensures TypeScript catches unknown providers at compile time
149+
// In practice, the discriminated union should prevent this at compile time
150+
const unknownProvider = "unknown" as any
151+
152+
await expect(
153+
getModels({
154+
provider: unknownProvider,
155+
}),
156+
).rejects.toThrow("Unknown provider: unknown")
157+
})
158+
})

0 commit comments

Comments
 (0)