mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-14 14:26:43 +00:00
feat: wire submit guardrail form to POST /guardrails/register
Backend: - Add required team_id to RegisterGuardrailRequest body - Validate team membership before allowing submission Frontend: - Add useRegisterGuardrail React Query mutation hook - Wire form onFinish to mutation, pass team_id from dropdown - Show notification and refresh submissions on success
This commit is contained in:
@@ -542,6 +542,7 @@ class RegisterGuardrailRequest(BaseModel):
|
||||
str, Any
|
||||
] # guardrail, mode, api_base required; api_key, headers, etc. optional
|
||||
guardrail_info: Optional[Dict[str, Any]] = None
|
||||
team_id: str
|
||||
|
||||
def get_litellm_params_dict(self) -> Dict[str, Any]:
|
||||
return dict(self.litellm_params)
|
||||
@@ -603,11 +604,29 @@ async def register_guardrail(
|
||||
if prisma_client is None:
|
||||
raise HTTPException(status_code=500, detail="Prisma client not initialized")
|
||||
|
||||
if not user_api_key_dict.team_id:
|
||||
if not request.team_id:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="Registration requires an API key associated with a team. Use a team-scoped key.",
|
||||
detail="team_id is required.",
|
||||
)
|
||||
team_id = request.team_id
|
||||
|
||||
# Validate the user is a member of the specified team
|
||||
if team_id != user_api_key_dict.team_id:
|
||||
from litellm.proxy.auth.auth_checks import get_team_membership
|
||||
from litellm.proxy.proxy_server import user_api_key_cache
|
||||
|
||||
membership = await get_team_membership(
|
||||
user_id=user_api_key_dict.user_id or "",
|
||||
team_id=request.team_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
)
|
||||
if membership is None:
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail=f"You are not a member of team {request.team_id!r}",
|
||||
)
|
||||
|
||||
params = request.get_litellm_params_dict()
|
||||
if params.get("guardrail") != GENERIC_GUARDRAIL_API:
|
||||
@@ -673,7 +692,7 @@ async def register_guardrail(
|
||||
"litellm_params": litellm_params_str,
|
||||
"guardrail_info": guardrail_info_str,
|
||||
"status": "pending_review",
|
||||
"team_id": user_api_key_dict.team_id,
|
||||
"team_id": team_id,
|
||||
"submitted_at": now,
|
||||
"created_at": now,
|
||||
"updated_at": now,
|
||||
|
||||
@@ -0,0 +1,74 @@
|
||||
import { useMutation, useQueryClient } from "@tanstack/react-query";
|
||||
import {
|
||||
getProxyBaseUrl,
|
||||
getGlobalLitellmHeaderName,
|
||||
deriveErrorMessage,
|
||||
handleError,
|
||||
} from "@/components/networking";
|
||||
import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
|
||||
import { createQueryKeys } from "../common/queryKeysFactory";
|
||||
|
||||
// ── Types ────────────────────────────────────────────────────────────────────
|
||||
|
||||
export interface RegisterGuardrailParams {
|
||||
guardrail_name: string;
|
||||
litellm_params: Record<string, unknown>;
|
||||
guardrail_info?: Record<string, unknown>;
|
||||
team_id: string;
|
||||
}
|
||||
|
||||
export interface RegisterGuardrailResponse {
|
||||
guardrail_id: string;
|
||||
guardrail_name: string;
|
||||
status: string;
|
||||
submitted_at?: string | null;
|
||||
}
|
||||
|
||||
// ── Fetch function ───────────────────────────────────────────────────────────
|
||||
|
||||
const registerGuardrail = async (
|
||||
accessToken: string,
|
||||
params: RegisterGuardrailParams,
|
||||
): Promise<RegisterGuardrailResponse> => {
|
||||
const baseUrl = getProxyBaseUrl();
|
||||
const url = `${baseUrl}/guardrails/register`;
|
||||
|
||||
const response = await fetch(url, {
|
||||
method: "POST",
|
||||
headers: {
|
||||
[getGlobalLitellmHeaderName()]: `Bearer ${accessToken}`,
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
body: JSON.stringify(params),
|
||||
});
|
||||
|
||||
if (!response.ok) {
|
||||
const errorData = await response.json().catch(() => ({}));
|
||||
const errorMessage = deriveErrorMessage(errorData);
|
||||
handleError(errorMessage);
|
||||
throw new Error(errorMessage);
|
||||
}
|
||||
|
||||
return response.json();
|
||||
};
|
||||
|
||||
// ── Hook ─────────────────────────────────────────────────────────────────────
|
||||
|
||||
const guardrailKeys = createQueryKeys("guardrails");
|
||||
|
||||
export const useRegisterGuardrail = () => {
|
||||
const { accessToken } = useAuthorized();
|
||||
const queryClient = useQueryClient();
|
||||
|
||||
return useMutation<RegisterGuardrailResponse, Error, RegisterGuardrailParams>({
|
||||
mutationFn: async (params) => {
|
||||
if (!accessToken) {
|
||||
throw new Error("Access token is required");
|
||||
}
|
||||
return registerGuardrail(accessToken, params);
|
||||
},
|
||||
onSuccess: () => {
|
||||
queryClient.invalidateQueries({ queryKey: guardrailKeys.all });
|
||||
},
|
||||
});
|
||||
};
|
||||
@@ -24,6 +24,7 @@ import {
|
||||
} from "@/components/networking";
|
||||
import NotificationsManager from "@/components/molecules/notifications_manager";
|
||||
import TeamDropdown from "@/components/common_components/team_dropdown";
|
||||
import { useRegisterGuardrail } from "@/app/(dashboard)/hooks/guardrails/useRegisterGuardrail";
|
||||
|
||||
type GuardrailStatus = "active" | "pending" | "rejected";
|
||||
|
||||
@@ -824,6 +825,7 @@ export function TeamGuardrailsTab({ accessToken }: TeamGuardrailsTabProps) {
|
||||
const [searchDebounced, setSearchDebounced] = useState("");
|
||||
const [isSubmitModalOpen, setIsSubmitModalOpen] = useState(false);
|
||||
const [submitForm] = Form.useForm();
|
||||
const registerGuardrail = useRegisterGuardrail();
|
||||
|
||||
useEffect(() => {
|
||||
const t = setTimeout(() => setSearchDebounced(search), 300);
|
||||
@@ -1099,22 +1101,27 @@ export function TeamGuardrailsTab({ accessToken }: TeamGuardrailsTabProps) {
|
||||
form={submitForm}
|
||||
layout="vertical"
|
||||
initialValues={{ mode: "pre_call" }}
|
||||
onFinish={(values) => {
|
||||
onFinish={async (values) => {
|
||||
const litellm_params: Record<string, unknown> = {
|
||||
guardrail: "generic_guardrail_api",
|
||||
mode: values.mode,
|
||||
api_base: values.api_base,
|
||||
...(values.extra_litellm_params ? JSON.parse(values.extra_litellm_params) : {}),
|
||||
};
|
||||
const payload = {
|
||||
guardrail_name: values.guardrail_name,
|
||||
litellm_params,
|
||||
guardrail_info: values.guardrail_info ? JSON.parse(values.guardrail_info) : undefined,
|
||||
};
|
||||
// TODO: call registerGuardrailCall once backend is wired
|
||||
console.log("Submit guardrail:", payload);
|
||||
setIsSubmitModalOpen(false);
|
||||
submitForm.resetFields();
|
||||
try {
|
||||
await registerGuardrail.mutateAsync({
|
||||
team_id: values.team_id,
|
||||
guardrail_name: values.guardrail_name,
|
||||
litellm_params,
|
||||
guardrail_info: values.guardrail_info ? JSON.parse(values.guardrail_info) : undefined,
|
||||
});
|
||||
NotificationsManager.success("Guardrail submitted for review");
|
||||
setIsSubmitModalOpen(false);
|
||||
submitForm.resetFields();
|
||||
fetchSubmissions();
|
||||
} catch {
|
||||
// error already handled by networking layer
|
||||
}
|
||||
}}
|
||||
>
|
||||
<Form.Item
|
||||
|
||||
Reference in New Issue
Block a user