Skip to content

Commit e093429

Browse files
yuneng-jiangyuneng-jiang
authored
Merge pull request #17177 from BerriAI/litellm_ui_model_perf_2
[Infra] UI - Migrate Provider Fields to React Query
2 parents 772be17 + 7e3f3c6 commit e093429

File tree

5 files changed

+353
-252
lines changed

5 files changed

+353
-252
lines changed

‎ui/litellm-dashboard/src/app/(dashboard)/hooks/providers/useProviderFields.ts‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
import { getProviderCreateMetadata, ProviderCreateInfo } from "@/components/networking";
2+
import { useQuery } from "@tanstack/react-query";
3+
import { createQueryKeys } from "../common/queryKeysFactory";
4+
5+
const providerFieldsKeys = createQueryKeys("providerFields");
6+
7+
export const useProviderFields = () => {
8+
return useQuery<ProviderCreateInfo[]>({
9+
queryKey: providerFieldsKeys.list({}),
10+
queryFn: async () => await getProviderCreateMetadata(),
11+
staleTime: 24 * 60 * 60 * 1000, // 24 hours - data rarely changes
12+
gcTime: 24 * 60 * 60 * 1000, // 24 hours - keep in cache for 24 hours
13+
});
14+
};

‎ui/litellm-dashboard/src/components/add_model/add_model_tab.test.tsx‎

Lines changed: 204 additions & 91 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,13 @@
1-
import { render, renderHook, waitFor } from "@testing-library/react";
2-
import { describe, it, vi, expect } from "vitest";
1+
import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
2+
import { render, renderHook, screen } from "@testing-library/react";
33
import { Form } from "antd";
4-
import AddModelTab from "./add_model_tab";
5-
import { Providers } from "../provider_info_helpers";
4+
import type { UploadProps } from "antd/es/upload";
5+
import { describe, expect, it, vi } from "vitest";
66
import type { Team } from "../key_team_helpers/key_list";
77
import type { CredentialItem } from "../networking";
8-
import type { UploadProps } from "antd/es/upload";
8+
import { Providers } from "../provider_info_helpers";
9+
import AddModelTab from "./add_model_tab";
910

10-
// Mock the networking module
1111
vi.mock("../networking", async () => {
1212
const actual = await vi.importActual("../networking");
1313
return {
@@ -19,6 +19,12 @@ vi.mock("../networking", async () => {
1919
modelAvailableCall: vi.fn().mockResolvedValue({
2020
data: [{ id: "model-group-1" }, { id: "model-group-2" }],
2121
}),
22+
modelHubCall: vi.fn().mockResolvedValue({
23+
data: [
24+
{ model_group: "gpt-4", mode: "chat" },
25+
{ model_group: "gpt-3.5-turbo", mode: "chat" },
26+
],
27+
}),
2228
getProviderCreateMetadata: vi.fn().mockResolvedValue([
2329
{
2430
provider: "OpenAI",
@@ -31,91 +37,198 @@ vi.mock("../networking", async () => {
3137
};
3238
});
3339

40+
const createQueryClient = () =>
41+
new QueryClient({
42+
defaultOptions: {
43+
queries: {
44+
retry: false,
45+
gcTime: 0,
46+
},
47+
},
48+
});
49+
50+
const createTestProps = () => {
51+
const { result } = renderHook(() => Form.useForm());
52+
const [form] = result.current;
53+
54+
const handleOk = vi.fn();
55+
const setSelectedProvider = vi.fn();
56+
const setProviderModelsFn = vi.fn();
57+
const getPlaceholder = vi.fn((provider: Providers) => `Enter ${provider} model name`);
58+
const setShowAdvancedSettings = vi.fn();
59+
60+
const selectedProvider = Providers.OpenAI;
61+
const providerModels = ["gpt-4", "gpt-3.5-turbo"];
62+
const showAdvancedSettings = false;
63+
64+
const teams: Team[] = [
65+
{
66+
team_id: "team-1",
67+
team_alias: "Test Team",
68+
models: ["gpt-4"],
69+
max_budget: 100,
70+
budget_duration: "monthly",
71+
tpm_limit: null,
72+
rpm_limit: null,
73+
organization_id: "org-1",
74+
created_at: "2024-01-01T00:00:00Z",
75+
keys: [],
76+
members_with_roles: [],
77+
},
78+
];
79+
80+
const credentials: CredentialItem[] = [
81+
{
82+
credential_name: "test-credential",
83+
credential_values: {},
84+
credential_info: {
85+
custom_llm_provider: "openai",
86+
description: "Test credential",
87+
},
88+
},
89+
];
90+
91+
const uploadProps: UploadProps = {
92+
beforeUpload: () => false,
93+
showUploadList: false,
94+
};
95+
96+
return {
97+
form,
98+
handleOk,
99+
setSelectedProvider,
100+
setProviderModelsFn,
101+
getPlaceholder,
102+
setShowAdvancedSettings,
103+
selectedProvider,
104+
providerModels,
105+
showAdvancedSettings,
106+
teams,
107+
credentials,
108+
uploadProps,
109+
accessToken: "test-access-token",
110+
userRole: "Admin",
111+
premiumUser: true,
112+
};
113+
};
114+
34115
describe("Add Model Tab", () => {
35-
it(
36-
"should render",
37-
async () => {
38-
// Create a form instance using renderHook
39-
const { result } = renderHook(() => Form.useForm());
40-
const [form] = result.current;
41-
42-
// Mock functions
43-
const handleOk = vi.fn();
44-
const setSelectedProvider = vi.fn();
45-
const setProviderModelsFn = vi.fn();
46-
const getPlaceholder = vi.fn((provider: Providers) => `Enter ${provider} model name`);
47-
const setShowAdvancedSettings = vi.fn();
48-
49-
// Mock data
50-
const selectedProvider = Providers.OpenAI;
51-
const providerModels = ["gpt-4", "gpt-3.5-turbo"];
52-
const showAdvancedSettings = false;
53-
54-
const teams: Team[] = [
55-
{
56-
team_id: "team-1",
57-
team_alias: "Test Team",
58-
models: ["gpt-4"],
59-
max_budget: 100,
60-
budget_duration: "monthly",
61-
tpm_limit: null,
62-
rpm_limit: null,
63-
organization_id: "org-1",
64-
created_at: "2024-01-01T00:00:00Z",
65-
keys: [],
66-
members_with_roles: [],
67-
},
68-
];
69-
70-
const credentials: CredentialItem[] = [
71-
{
72-
credential_name: "test-credential",
73-
credential_values: {},
74-
credential_info: {
75-
custom_llm_provider: "openai",
76-
description: "Test credential",
77-
},
78-
},
79-
];
80-
81-
const uploadProps: UploadProps = {
82-
beforeUpload: () => false,
83-
showUploadList: false,
84-
};
85-
86-
const accessToken = "test-access-token";
87-
const userRole = "Admin";
88-
const premiumUser = true;
89-
90-
const { container, findByText } = render(
116+
it("should render", async () => {
117+
const props = createTestProps();
118+
const queryClient = createQueryClient();
119+
120+
render(
121+
<QueryClientProvider client={queryClient}>
91122
<AddModelTab
92-
form={form}
93-
handleOk={handleOk}
94-
selectedProvider={selectedProvider}
95-
setSelectedProvider={setSelectedProvider}
96-
providerModels={providerModels}
97-
setProviderModelsFn={setProviderModelsFn}
98-
getPlaceholder={getPlaceholder}
99-
uploadProps={uploadProps}
100-
showAdvancedSettings={showAdvancedSettings}
101-
setShowAdvancedSettings={setShowAdvancedSettings}
102-
teams={teams}
103-
credentials={credentials}
104-
accessToken={accessToken}
105-
userRole={userRole}
106-
premiumUser={premiumUser}
107-
/>,
108-
);
109-
110-
// Wait for the tabs to render which indicates the component loaded
111-
await waitFor(
112-
() => {
113-
const tabs = container.querySelectorAll('[role="tab"]');
114-
expect(tabs.length).toBeGreaterThan(0);
115-
},
116-
{ timeout: 10000 },
117-
);
118-
},
119-
15000,
120-
);
123+
form={props.form}
124+
handleOk={props.handleOk}
125+
selectedProvider={props.selectedProvider}
126+
setSelectedProvider={props.setSelectedProvider}
127+
providerModels={props.providerModels}
128+
setProviderModelsFn={props.setProviderModelsFn}
129+
getPlaceholder={props.getPlaceholder}
130+
uploadProps={props.uploadProps}
131+
showAdvancedSettings={props.showAdvancedSettings}
132+
setShowAdvancedSettings={props.setShowAdvancedSettings}
133+
teams={props.teams}
134+
credentials={props.credentials}
135+
accessToken={props.accessToken}
136+
userRole={props.userRole}
137+
premiumUser={props.premiumUser}
138+
/>
139+
</QueryClientProvider>,
140+
);
141+
142+
expect(await screen.findByRole("tab", { name: "Add Model" })).toBeInTheDocument();
143+
});
144+
145+
it("should display both Add Model and Add Auto Router tabs", async () => {
146+
const props = createTestProps();
147+
const queryClient = createQueryClient();
148+
149+
render(
150+
<QueryClientProvider client={queryClient}>
151+
<AddModelTab
152+
form={props.form}
153+
handleOk={props.handleOk}
154+
selectedProvider={props.selectedProvider}
155+
setSelectedProvider={props.setSelectedProvider}
156+
providerModels={props.providerModels}
157+
setProviderModelsFn={props.setProviderModelsFn}
158+
getPlaceholder={props.getPlaceholder}
159+
uploadProps={props.uploadProps}
160+
showAdvancedSettings={props.showAdvancedSettings}
161+
setShowAdvancedSettings={props.setShowAdvancedSettings}
162+
teams={props.teams}
163+
credentials={props.credentials}
164+
accessToken={props.accessToken}
165+
userRole={props.userRole}
166+
premiumUser={props.premiumUser}
167+
/>
168+
</QueryClientProvider>,
169+
);
170+
171+
expect(await screen.findByRole("tab", { name: "Add Model" })).toBeInTheDocument();
172+
expect(await screen.findByRole("tab", { name: "Add Auto Router" })).toBeInTheDocument();
173+
});
174+
175+
it("should display provider selection field", async () => {
176+
const props = createTestProps();
177+
const queryClient = createQueryClient();
178+
179+
render(
180+
<QueryClientProvider client={queryClient}>
181+
<AddModelTab
182+
form={props.form}
183+
handleOk={props.handleOk}
184+
selectedProvider={props.selectedProvider}
185+
setSelectedProvider={props.setSelectedProvider}
186+
providerModels={props.providerModels}
187+
setProviderModelsFn={props.setProviderModelsFn}
188+
getPlaceholder={props.getPlaceholder}
189+
uploadProps={props.uploadProps}
190+
showAdvancedSettings={props.showAdvancedSettings}
191+
setShowAdvancedSettings={props.setShowAdvancedSettings}
192+
teams={props.teams}
193+
credentials={props.credentials}
194+
accessToken={props.accessToken}
195+
userRole={props.userRole}
196+
premiumUser={props.premiumUser}
197+
/>
198+
</QueryClientProvider>,
199+
);
200+
201+
expect(await screen.findByText("Provider")).toBeInTheDocument();
202+
});
203+
204+
it("should display Test Connect and Add Model buttons", async () => {
205+
const props = createTestProps();
206+
const queryClient = createQueryClient();
207+
208+
render(
209+
<QueryClientProvider client={queryClient}>
210+
<AddModelTab
211+
form={props.form}
212+
handleOk={props.handleOk}
213+
selectedProvider={props.selectedProvider}
214+
setSelectedProvider={props.setSelectedProvider}
215+
providerModels={props.providerModels}
216+
setProviderModelsFn={props.setProviderModelsFn}
217+
getPlaceholder={props.getPlaceholder}
218+
uploadProps={props.uploadProps}
219+
showAdvancedSettings={props.showAdvancedSettings}
220+
setShowAdvancedSettings={props.setShowAdvancedSettings}
221+
teams={props.teams}
222+
credentials={props.credentials}
223+
accessToken={props.accessToken}
224+
userRole={props.userRole}
225+
premiumUser={props.premiumUser}
226+
/>
227+
</QueryClientProvider>,
228+
);
229+
230+
const testConnectButtons = await screen.findAllByRole("button", { name: "Test Connect" });
231+
expect(testConnectButtons.length).toBeGreaterThan(0);
232+
expect(await screen.findByRole("button", { name: "Add Model" })).toBeInTheDocument();
233+
}, 10000); // 10 seconds timeout for complex logic
121234
});

0 commit comments

Comments
 (0)