mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
refactor(ui/guardrails): expose full data from useGuardrails hook
Extend useGuardrails to return the full guardrail objects plus derived globalGuardrailNames / optionalGuardrailNames sets via React Query's select option, instead of just an array of names. Update its existing consumer (AddModelForm) to extract names from the new shape. The previous shape was tailored to AddModelForm's single use case (populate a Select with names). The team info per-guardrail opt-out work needs default_on per guardrail to split globals from non-globals, which the old shape couldn't provide. Consolidating into the existing hook gives both consumers one source of truth and one React Query cache entry instead of two parallel fetches. - useGuardrails.ts: rewrite return type, derive global/optional sets in select(); preserve the existing query key and auth-gate semantics - AddModelForm.tsx: extract names from data?.guardrails.map(...) - AddModelForm.test.tsx: update mock to return the new shape (also fixes a pre-existing shape mismatch in the mock) - useGuardrails.test.ts: update 3 assertions to read names via data?.guardrails.map(...) instead of asserting against the flat array
This commit is contained in:
parent
88beed905e
commit
a9d64f8620
4 changed files with 56 additions and 14 deletions
|
|
@ -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");
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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<string, unknown> | null;
|
||||
guardrail_id?: string | null;
|
||||
[key: string]: unknown;
|
||||
}
|
||||
|
||||
interface GuardrailsListResponse {
|
||||
guardrails: GuardrailListItem[];
|
||||
}
|
||||
|
||||
export interface GuardrailsListData {
|
||||
guardrails: GuardrailListItem[];
|
||||
globalGuardrailNames: Set<string>;
|
||||
optionalGuardrailNames: Set<string>;
|
||||
}
|
||||
|
||||
// ── Hook ─────────────────────────────────────────────────────────────────────
|
||||
|
||||
const guardrailKeys = createQueryKeys("guardrails");
|
||||
|
||||
export const useGuardrails = (): UseQueryResult<string[]> => {
|
||||
export const useGuardrails = (): UseQueryResult<GuardrailsListData> => {
|
||||
const { accessToken, userId, userRole } = useAuthorized();
|
||||
return useQuery<string[]>({
|
||||
return useQuery<GuardrailsListResponse, Error, GuardrailsListData>({
|
||||
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<string>();
|
||||
const optionalGuardrailNames = new Set<string>();
|
||||
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 };
|
||||
},
|
||||
});
|
||||
};
|
||||
|
|
|
|||
|
|
@ -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<string>(),
|
||||
optionalGuardrailNames: new Set<string>(["test-guardrail"]),
|
||||
},
|
||||
isLoading: false,
|
||||
error: null,
|
||||
}),
|
||||
|
|
|
|||
|
|
@ -63,7 +63,8 @@ const AddModelForm: React.FC<AddModelFormProps> = ({
|
|||
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 () => {
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue