diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/guardrails/useGuardrails.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/guardrails/useGuardrails.test.ts index d9e96a5308c..d5b788ab6a5 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/guardrails/useGuardrails.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/guardrails/useGuardrails.test.ts @@ -74,7 +74,7 @@ describe("useGuardrails", () => { expect(result.current.isSuccess).toBe(true); }); - expect(result.current.data).toEqual(expectedGuardrailNames); + expect(result.current.data?.guardrails.map((g) => g.guardrail_name)).toEqual(expectedGuardrailNames); expect(result.current.error).toBeNull(); expect(getGuardrailsList).toHaveBeenCalledWith("test-access-token"); expect(getGuardrailsList).toHaveBeenCalledTimes(1); @@ -228,7 +228,7 @@ describe("useGuardrails", () => { expect(result.current.isSuccess).toBe(true); }); - expect(result.current.data).toEqual([]); + expect(result.current.data?.guardrails).toEqual([]); expect(getGuardrailsList).toHaveBeenCalledWith("test-access-token"); }); @@ -265,9 +265,10 @@ describe("useGuardrails", () => { expect(result.current.isSuccess).toBe(true); }); - expect(result.current.data).toEqual(expectedNames); - expect(result.current.data).toHaveLength(2); - expect(result.current.data).toContain("custom-guardrail-1"); - expect(result.current.data).toContain("custom-guardrail-2"); + const names = result.current.data?.guardrails.map((g) => g.guardrail_name); + expect(names).toEqual(expectedNames); + expect(names).toHaveLength(2); + expect(names).toContain("custom-guardrail-1"); + expect(names).toContain("custom-guardrail-2"); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/guardrails/useGuardrails.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/guardrails/useGuardrails.ts index 9786b7fa359..5c7a8df050b 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/guardrails/useGuardrails.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/guardrails/useGuardrails.ts @@ -3,16 +3,52 @@ import { createQueryKeys } from "../common/queryKeysFactory"; import { getGuardrailsList } from "@/components/networking"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +// ── Types ──────────────────────────────────────────────────────────────────── + +export interface GuardrailListItem { + guardrail_name: string; + litellm_params?: { + default_on?: boolean; + mode?: string | string[]; + [key: string]: unknown; + }; + guardrail_info?: Record | null; + guardrail_id?: string | null; + [key: string]: unknown; +} + +interface GuardrailsListResponse { + guardrails: GuardrailListItem[]; +} + +export interface GuardrailsListData { + guardrails: GuardrailListItem[]; + globalGuardrailNames: Set; + optionalGuardrailNames: Set; +} + +// ── Hook ───────────────────────────────────────────────────────────────────── + const guardrailKeys = createQueryKeys("guardrails"); -export const useGuardrails = (): UseQueryResult => { +export const useGuardrails = (): UseQueryResult => { const { accessToken, userId, userRole } = useAuthorized(); - return useQuery({ + return useQuery({ queryKey: guardrailKeys.list({}), - queryFn: async () => { - const response = await getGuardrailsList(accessToken!); - return response.guardrails.map((g: { guardrail_name: string }) => g.guardrail_name); - }, + queryFn: async () => getGuardrailsList(accessToken!), enabled: Boolean(accessToken && userId && userRole), + select: (data) => { + const guardrails: GuardrailListItem[] = data?.guardrails ?? []; + const globalGuardrailNames = new Set(); + const optionalGuardrailNames = new Set(); + for (const g of guardrails) { + if (g.litellm_params?.default_on) { + globalGuardrailNames.add(g.guardrail_name); + } else { + optionalGuardrailNames.add(g.guardrail_name); + } + } + return { guardrails, globalGuardrailNames, optionalGuardrailNames }; + }, }); }; diff --git a/ui/litellm-dashboard/src/components/add_model/AddModelForm.test.tsx b/ui/litellm-dashboard/src/components/add_model/AddModelForm.test.tsx index de5e1387396..4a3dbaedf74 100644 --- a/ui/litellm-dashboard/src/components/add_model/AddModelForm.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/AddModelForm.test.tsx @@ -82,7 +82,11 @@ vi.mock("@/app/(dashboard)/hooks/teams/useTeams", () => ({ vi.mock("@/app/(dashboard)/hooks/guardrails/useGuardrails", () => ({ useGuardrails: vi.fn().mockReturnValue({ - data: [{ guardrail_name: "test-guardrail" }], + data: { + guardrails: [{ guardrail_name: "test-guardrail" }], + globalGuardrailNames: new Set(), + optionalGuardrailNames: new Set(["test-guardrail"]), + }, isLoading: false, error: null, }), diff --git a/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx b/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx index 65e239d58ea..cb35a6f4142 100644 --- a/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx +++ b/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx @@ -63,7 +63,8 @@ const AddModelForm: React.FC = ({ isLoading: isProviderMetadataLoading, error: providerMetadataError, } = useProviderFields(); - const { data: guardrailsList, isLoading: isGuardrailsLoading, error: guardrailsError } = useGuardrails(); + const { data: guardrailsData, isLoading: isGuardrailsLoading, error: guardrailsError } = useGuardrails(); + const guardrailsList = guardrailsData?.guardrails.map((g) => g.guardrail_name); const { data: tagsList, isLoading: isTagsLoading, error: tagsError } = useTags(); const handleTestConnection = async () => {