ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Refactor react-query hooks

Browse Source
This commit is contained in:
yuneng-jiang 2025-12-24 10:13:55 -08:00
parent 8a3d4967ed
commit 825503fa1b
18 changed files with 63 additions and 47 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
ui/litellm-dashboard/src/app/(dashboard)/hooks/agents/useAgents.ts
View File

@ -3,10 +3,12 @@ import { AgentsResponse } from "@/components/agents/types";
import { useQuery } from "@tanstack/react-query";
import { createQueryKeys } from "../common/queryKeysFactory";
import { all_admin_roles } from "@/utils/roles";
import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
const agentsKeys = createQueryKeys("agents");
export const useAgents = (accessToken: string | null, userRole: string | null) => {
export const useAgents = () => {
const { accessToken, userRole } = useAuthorized();
return useQuery<AgentsResponse>({
queryKey: agentsKeys.list({}),
queryFn: async () => await getAgentsList(accessToken!),

4
ui/litellm-dashboard/src/app/(dashboard)/hooks/credentials/useCredentials.ts
View File

@ -1,10 +1,12 @@
import { credentialListCall, CredentialsResponse } from "@/components/networking";
import { useQuery } from "@tanstack/react-query";
import { createQueryKeys } from "../common/queryKeysFactory";
import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
const credentialsKeys = createQueryKeys("credentials");
export const useCredentials = (accessToken: string | null) => {
export const useCredentials = () => {
const { accessToken } = useAuthorized();
return useQuery<CredentialsResponse>({
queryKey: credentialsKeys.list({}),
queryFn: async () => await credentialListCall(accessToken!),

7
ui/litellm-dashboard/src/app/(dashboard)/hooks/customers/useCustomers.ts
View File

@ -2,7 +2,7 @@ import { allEndUsersCall } from "@/components/networking";
import { useQuery } from "@tanstack/react-query";
import { createQueryKeys } from "../common/queryKeysFactory";
import { all_admin_roles } from "@/utils/roles";
import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
const customersKeys = createQueryKeys("customers");
export interface Customer {
@ -32,10 +32,11 @@ export interface Customer {
export type CustomersResponse = Customer[];
export const useCustomers = (accessToken: string | null, userRole: string | null) => {
export const useCustomers = () => {
const { accessToken, userRole } = useAuthorized();
return useQuery<CustomersResponse>({
queryKey: customersKeys.list({}),
queryFn: async () => await allEndUsersCall(accessToken!),
enabled: Boolean(accessToken) && all_admin_roles.includes(userRole || ""),
enabled: Boolean(accessToken) && all_admin_roles.includes(userRole!),
});
};

7
ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPAccessGroups.ts
View File

@ -1,13 +1,14 @@
import { useQuery } from "@tanstack/react-query";
import { createQueryKeys } from "../common/queryKeysFactory";
import { fetchMCPAccessGroups } from "@/components/networking";
import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
const mcpAccessGroupsKeys = createQueryKeys("mcpAccessGroups");
export const useMCPAccessGroups = (accessToken: string | null) => {
export const useMCPAccessGroups = () => {
const { accessToken } = useAuthorized();
return useQuery<string[]>({
queryKey: mcpAccessGroupsKeys.list({}),
queryFn: async () => await fetchMCPAccessGroups(accessToken!),
enabled: !!accessToken,
enabled: Boolean(accessToken),
});
};

4
ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServers.ts
View File

@ -2,10 +2,12 @@ import { useQuery } from "@tanstack/react-query";
import { createQueryKeys } from "../common/queryKeysFactory";
import { fetchMCPServers } from "@/components/networking";
import { MCPServer } from "@/components/mcp_tools/types";
import useAuthorized from "../useAuthorized";
const mcpServersKeys = createQueryKeys("mcpServers");
export const useMCPServers = (accessToken: string | null) => {
export const useMCPServers = () => {
const { accessToken } = useAuthorized();
return useQuery<MCPServer[]>({
queryKey: mcpServersKeys.list({}),
queryFn: async () => await fetchMCPServers(accessToken!),

14
ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts
View File

@ -1,24 +1,26 @@
import { useQuery } from "@tanstack/react-query";
import { createQueryKeys } from "../common/queryKeysFactory";
import { modelInfoCall, modelHubCall } from "@/components/networking";
import useAuthorized from "../useAuthorized";
const modelKeys = createQueryKeys("models");
const modelHubKeys = createQueryKeys("modelHub");
export const useModelsInfo = (accessToken: string | null, userID: string | null, userRole: string | null) => {
export const useModelsInfo = () => {
const { accessToken, userId, userRole } = useAuthorized();
return useQuery({
queryKey: modelKeys.list({
filters: {
...(userID && { userID }),
...(userId && { userId }),
...(userRole && { userRole }),
},
}),
queryFn: async () => await modelInfoCall(accessToken!, userID!, userRole!),
enabled: Boolean(accessToken && userID && userRole),
queryFn: async () => await modelInfoCall(accessToken!, userId!, userRole!),
enabled: Boolean(accessToken && userId && userRole),
});
};
export const useModelHub = (accessToken: string | null) => {
export const useModelHub = () => {
const { accessToken } = useAuthorized();
return useQuery({
queryKey: modelHubKeys.list({}),
queryFn: async () => await modelHubCall(accessToken!),

7
ui/litellm-dashboard/src/app/(dashboard)/hooks/organizations/useOrganizations.ts
View File

@ -1,13 +1,16 @@
import { useQuery, UseQueryResult } from "@tanstack/react-query";
import { createQueryKeys } from "../common/queryKeysFactory";
import { organizationListCall, Organization } from "@/components/networking";
import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
const organizationKeys = createQueryKeys("organizations");
export const useOrganizations = (accessToken: string | null): UseQueryResult<Organization[]> => {
export const useOrganizations = (): UseQueryResult<Organization[]> => {
const { accessToken } = useAuthorized();
const { userId, userRole } = useAuthorized();
return useQuery<Organization[]>({
queryKey: organizationKeys.list({}),
queryFn: async () => await organizationListCall(accessToken!),
enabled: Boolean(accessToken),
enabled: Boolean(accessToken && userId && userRole),
});
};

4
ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.ts
View File

@ -7,11 +7,11 @@ import { createQueryKeys } from "@/app/(dashboard)/hooks/common/queryKeysFactory
const teamKeys = createQueryKeys("teams");
export const useTeams = (): UseQueryResult<Team[]> => {
const { accessToken, userId: userID, userRole } = useAuthorized();
const { accessToken, userId, userRole } = useAuthorized();
return useQuery<Team[]>({
queryKey: teamKeys.list({}),
queryFn: async () => await fetchTeams(accessToken!, userID, userRole, null),
queryFn: async () => await fetchTeams(accessToken!, userId, userRole, null),
enabled: Boolean(accessToken),
});
};

3
ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.test.ts
View File

@ -5,6 +5,9 @@ import { afterEach, describe, expect, it, vi } from "vitest";
import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
import useAuthorized from "./useAuthorized";
// Unmock useAuthorized to test the actual implementation
vi.unmock("@/app/(dashboard)/hooks/useAuthorized");
const { replaceMock, clearTokenCookiesMock, getProxyBaseUrlMock, getUiConfigMock } = vi.hoisted(() => ({
replaceMock: vi.fn(),
clearTokenCookiesMock: vi.fn(),

13
ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.test.tsx
View File

@ -37,19 +37,6 @@ vi.mock("@/app/(dashboard)/models-and-endpoints/components/ModelAnalyticsTab/Mod
default: () => null,
}));
vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({
default: () => ({
token: "123",
accessToken: "123",
userId: "user-1",
userEmail: "user@example.com",
userRole: "Admin",
premiumUser: false,
disabledPersonalKeyCreation: null,
showSSOBanner: false,
}),
}));
vi.mock("@/app/(dashboard)/hooks/useTeams", () => ({
default: () => ({
teams: [],

8
ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx
View File

@ -152,12 +152,8 @@ const ModelsAndEndpointsView: React.FC<ModelDashboardProps> = ({
const [selectedTabIndex, setSelectedTabIndex] = useState(0);
const queryClient = useQueryClient();
const {
data: modelDataResponse,
isLoading: isLoadingModels,
refetch: refetchModels,
} = useModelsInfo(accessToken, userID, userRole);
const { data: credentialsResponse } = useCredentials(accessToken);
const { data: modelDataResponse, isLoading: isLoadingModels, refetch: refetchModels } = useModelsInfo();
const { data: credentialsResponse } = useCredentials();
const credentialsList = credentialsResponse?.credentials || [];
const { data: uiSettings } = useUISettings(accessToken || "");

4
ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.tsx
View File

@ -80,8 +80,8 @@ const UsagePage: React.FC<UsagePageProps> = ({ teams, organizations }) => {
});
const [allTags, setAllTags] = useState<EntityList[]>([]);
const { data: customers = [] } = useCustomers(accessToken, userRole);
const { data: agentsResponse } = useAgents(accessToken, userRole);
const { data: customers = [] } = useCustomers();
const { data: agentsResponse } = useAgents();
const [modelViewType, setModelViewType] = useState<"groups" | "individual">("groups");
const [isCloudZeroModalOpen, setIsCloudZeroModalOpen] = useState(false);
const [isGlobalExportModalOpen, setIsGlobalExportModalOpen] = useState(false);

2
ui/litellm-dashboard/src/components/leftnav.tsx
View File

@ -57,7 +57,7 @@ interface MenuGroup {
const Sidebar: React.FC<SidebarProps> = ({ setPage, defaultSelectedKey, collapsed = false }) => {
const { userId, accessToken, userRole } = useAuthorized();
const { data: organizations } = useOrganizations(accessToken);
const { data: organizations } = useOrganizations();
// Check if user is an org_admin
const isOrgAdmin = useMemo(() => {

4
ui/litellm-dashboard/src/components/mcp_server_management/MCPServerSelector.tsx
View File

@ -23,8 +23,8 @@ const MCPServerSelector: React.FC<MCPServerSelectorProps> = ({
placeholder = "Select MCP servers",
disabled = false,
}) => {
const { data: mcpServers = [], isLoading: serversLoading } = useMCPServers(accessToken);
const { data: accessGroups = [], isLoading: groupsLoading } = useMCPAccessGroups(accessToken);
const { data: mcpServers = [], isLoading: serversLoading } = useMCPServers();
const { data: accessGroups = [], isLoading: groupsLoading } = useMCPAccessGroups();
const loading = serversLoading || groupsLoading;

4
ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.test.tsx
View File

@ -71,7 +71,9 @@ describe("MCPToolPermissions", () => {
});
// Verify API calls
expect(networking.fetchMCPServers).toHaveBeenCalledWith(mockAccessToken);
// Note: useMCPServers uses useAuthorized() internally, which returns "123" from global mock
expect(networking.fetchMCPServers).toHaveBeenCalledWith("123");
// listMCPTools uses the accessToken prop directly
expect(networking.listMCPTools).toHaveBeenCalledWith(mockAccessToken, mockServerId);
});

5
ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.test.tsx
View File

@ -32,7 +32,7 @@ const createQueryClient = () =>
describe("MCPServers", () => {
const defaultProps = {
accessToken: "test-token",
accessToken: "123",
userRole: "Admin",
userID: "admin-user-id",
};
@ -120,6 +120,7 @@ describe("MCPServers", () => {
expect(getByText("test-server-2")).toBeInTheDocument();
// Verify the API was called
expect(networking.fetchMCPServers).toHaveBeenCalledWith("test-token");
// Note: useMCPServers uses useAuthorized() internally, which returns "123" from global mock
expect(networking.fetchMCPServers).toHaveBeenCalledWith("123");
});
});

2
ui/litellm-dashboard/src/components/model_add/credentials.tsx
View File

@ -32,7 +32,7 @@ interface CredentialsPanelProps {
const CredentialsPanel: React.FC<CredentialsPanelProps> = ({ uploadProps }) => {
const { accessToken } = useAuthorized();
const { data: credentialsResponse, refetch: refetchCredentials } = useCredentials(accessToken);
const { data: credentialsResponse, refetch: refetchCredentials } = useCredentials();
const credentialList = credentialsResponse?.credentials || [];
const [isAddModalOpen, setIsAddModalOpen] = useState(false);

14
ui/litellm-dashboard/tests/setupTests.ts
View File

@ -33,6 +33,20 @@ vi.mock("@tremor/react", async (importOriginal) => {
};
});
// Global mock for useAuthorized hook to avoid repeating the same mock in every test file
vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({
default: () => ({
token: "123",
accessToken: "123",
userId: "user-1",
userEmail: "user@example.com",
userRole: "Admin",
premiumUser: false,
disabledPersonalKeyCreation: null,
showSSOBanner: false,
}),
}));
afterEach(() => {
cleanup();
});