phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-27 03:34:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

(UI) - Set/edit guardrails on a virtual key (#7954)

Browse source 
* Revert "JWT Auth - `enforce_rbac` support + UI team view, spend calc fix (#7863)"

This reverts commit dca6904937.

* Revert "Litellm dev 01 10 2025 p2 (#7679)"

This reverts commit c4780479a9.

* ui - allow setting guardrails on a key

* working edit guardrails

* fix edit guardrails on a key

* Revert "Revert "JWT Auth - `enforce_rbac` support + UI team view, spend calc fix (#7863)""

This reverts commit 8f7b9ae1af.

* Revert "Revert "Litellm dev 01 10 2025 p2 (#7679)""

This reverts commit a609139dde.

* fix edit guardrail on ui

* fix list_guardrails
This commit is contained in:
Ishaan Jaff 2025-01-23 18:01:54 -08:00 • committed by GitHub
parent e6ec4f21e5
commit 84fb5aead8
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 185 additions and 35 deletions
Download patch file Download diff file Expand all files Collapse all files

2
litellm/proxy/_types.py
View file

@ -2288,7 +2288,6 @@ class ProxyStateVariables(TypedDict):
UI_TEAM_ID = "litellm-dashboard"
class JWTAuthBuilderResult(TypedDict):
is_proxy_admin: bool
team_object: Optional[LiteLLM_TeamTable]
@ -2301,6 +2300,7 @@ class JWTAuthBuilderResult(TypedDict):
end_user_id: Optional[str]
org_id: Optional[str]
class ClientSideFallbackModel(TypedDict, total=False):
"""
Dictionary passed when client configuring input

12
litellm/proxy/guardrails/guardrail_endpoints.py
View file

@ -6,7 +6,6 @@ from typing import Dict, List, Optional, cast
from fastapi import APIRouter, Depends, HTTPException, status
from litellm.proxy._types import CommonProxyErrors
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.types.guardrails import GuardrailInfoResponse, ListGuardrailsResponse
@ -40,7 +39,6 @@ def _get_guardrails_list_response(
)
async def list_guardrails():
"""
✨ Enterprise Feature
List the guardrails that are available on the proxy server
👉 [Guardrail docs](https://docs.litellm.ai/docs/proxy/guardrails/quick_start)
@ -74,15 +72,7 @@ async def list_guardrails():
}
```
"""
from litellm.proxy.proxy_server import premium_user, proxy_config
if not premium_user:
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail={
"error": CommonProxyErrors.not_premium_user.value,
},
)
from litellm.proxy.proxy_server import proxy_config
config = proxy_config.config

16
litellm/proxy/proxy_config.yaml
View file

@ -14,6 +14,20 @@ model_list:
api_base: https://exampleopenaiendpoint-production.up.railway.app/
litellm_settings:
callbacks: ["prometheus"]
guardrails:
- guardrail_name: "bedrock-pre-guard"
litellm_params:
guardrail: bedrock # supported values: "aporia", "bedrock", "lakera"
mode: "during_call"
guardrailIdentifier: ff6ujrregl1q
guardrailVersion: "DRAFT"
- guardrail_name: "bedrock-post-guard"
litellm_params:
guardrail: bedrock # supported values: "aporia", "bedrock", "lakera"
mode: "post_call"
guardrailIdentifier: ff6ujrregl1q
guardrailVersion: "DRAFT"

48
ui/litellm-dashboard/src/components/create_key_button.tsx
View file

@ -27,7 +27,10 @@ import {
keyCreateCall,
slackBudgetAlertsHealthCheck,
modelAvailableCall,
getGuardrailsList,
} from "./networking";
import { InfoCircleOutlined } from '@ant-design/icons';
import { Tooltip } from 'antd';
const { Option } = Select;
@ -81,7 +84,7 @@ const CreateKey: React.FC<CreateKeyProps> = ({
const [modelsToPick, setModelsToPick] = useState([]);
const [keyOwner, setKeyOwner] = useState("you");
const [predefinedTags, setPredefinedTags] = useState(getPredefinedTags(data));
const [guardrailsList, setGuardrailsList] = useState<string[]>([]);
const handleOk = () => {
setIsModalVisible(false);
@ -121,6 +124,22 @@ const CreateKey: React.FC<CreateKeyProps> = ({
fetchUserModels();
}, [accessToken, userID, userRole]);
useEffect(() => {
const fetchGuardrails = async () => {
try {
const response = await getGuardrailsList(accessToken);
const guardrailNames = response.guardrails.map(
(g: { guardrail_name: string }) => g.guardrail_name
);
setGuardrailsList(guardrailNames);
} catch (error) {
console.error("Failed to fetch guardrails:", error);
}
};
fetchGuardrails();
}, [accessToken]);
const handleCreate = async (formValues: Record<string, any>) => {
try {
const newKeyAlias = formValues?.key_alias ?? "";
@ -392,6 +411,33 @@ const CreateKey: React.FC<CreateKeyProps> = ({
>
<TextInput placeholder="" />
</Form.Item>
<Form.Item
label={
<span>
Guardrails{' '}
<Tooltip title="Setup your first guardrail">
<a
href="https://docs.litellm.ai/docs/proxy/guardrails/quick_start"
target="_blank"
rel="noopener noreferrer"
onClick={(e) => e.stopPropagation()} // Prevent accordion from collapsing when clicking link
>
<InfoCircleOutlined style={{ marginLeft: '4px' }} />
</a>
</Tooltip>
</span>
}
name="guardrails"
className="mt-8"
help="Select existing guardrails or enter new ones"
>
<Select
mode="tags"
style={{ width: '100%' }}
placeholder="Select or enter guardrails"
options={guardrailsList.map(name => ({ value: name, label: name }))}
/>
</Form.Item>
<Form.Item label="Metadata" name="metadata" className="mt-8">
<Input.TextArea

28
ui/litellm-dashboard/src/components/networking.tsx
View file

@ -3050,3 +3050,31 @@ export const getProxyUISettings = async (
}
};
export const getGuardrailsList = async (accessToken: String) => {
try {
let url = proxyBaseUrl ? `${proxyBaseUrl}/guardrails/list` : `/guardrails/list`;
const response = await fetch(url, {
method: "GET",
headers: {
[globalLitellmHeaderName]: `Bearer ${accessToken}`,
"Content-Type": "application/json",
},
});
if (!response.ok) {
const errorData = await response.text();
handleError(errorData);
throw new Error("Network response was not ok");
}
const data = await response.json();
console.log("Guardrails list response:", data);
return data;
} catch (error) {
console.error("Failed to fetch guardrails list:", error);
throw error;
}
};

1
ui/litellm-dashboard/src/components/user_dashboard.tsx
View file

@ -206,7 +206,6 @@ const UserDashboard: React.FC<UserDashboardProps> = ({
setUserSpendData(response["user_info"]);
console.log(`userSpendData: ${JSON.stringify(userSpendData)}`)
const teamsArray = [...response["teams"]];
if (teamsArray.length > 0) {
console.log(`response['teams']: ${JSON.stringify(teamsArray)}`);

87
ui/litellm-dashboard/src/components/view_key_table.tsx
View file

@ -1,6 +1,6 @@
"use client";
import React, { useEffect, useState } from "react";
import { keyDeleteCall, modelAvailableCall } from "./networking";
import { keyDeleteCall, modelAvailableCall, getGuardrailsList } from "./networking";
import { add } from 'date-fns';
import { InformationCircleIcon, StatusOnlineIcon, TrashIcon, PencilAltIcon, RefreshIcon } from "@heroicons/react/outline";
import { keySpendLogsCall, PredictedSpendLogsCall, keyUpdateCall, modelInfoCall, regenerateKeyCall } from "./networking";
@ -26,6 +26,7 @@ import {
TextInput,
Textarea,
} from "@tremor/react";
import { InfoCircleOutlined } from '@ant-design/icons';
import { Select as Select3, SelectItem, MultiSelect, MultiSelectItem } from "@tremor/react";
import {
Button as Button2,
@ -130,6 +131,7 @@ const ViewKeyTable: React.FC<ViewKeyTableProps> = ({
const [newExpiryTime, setNewExpiryTime] = useState<string | null>(null);
const [knownTeamIDs, setKnownTeamIDs] = useState(initialKnownTeamIDs);
const [guardrailsList, setGuardrailsList] = useState<string[]>([]);
// Function to check if user is admin of a team
const isUserTeamAdmin = (team: any) => {
@ -311,22 +313,49 @@ const ViewKeyTable: React.FC<ViewKeyTableProps> = ({
const [keyTeam, setKeyTeam] = useState(selectedTeam);
const [errorModels, setErrorModels] = useState<string[]>([]);
const [errorBudget, setErrorBudget] = useState<boolean>(false);
const [guardrailsList, setGuardrailsList] = useState<string[]>([]);
useEffect(() => {
const fetchGuardrails = async () => {
try {
const response = await getGuardrailsList(accessToken);
const guardrailNames = response.guardrails.map(
(g: { guardrail_name: string }) => g.guardrail_name
);
setGuardrailsList(guardrailNames);
} catch (error) {
console.error("Failed to fetch guardrails:", error);
}
};
fetchGuardrails();
}, [accessToken]);
let metadataString = '';
try {
metadataString = JSON.stringify(token.metadata, null, 2);
// Create a copy of metadata without guardrails for display
const displayMetadata = { ...token.metadata };
delete displayMetadata.guardrails;
metadataString = JSON.stringify(displayMetadata, null, 2);
} catch (error) {
console.error("Error stringifying metadata:", error);
// You can choose a fallback, such as an empty string or a warning message
metadataString = '';
}
// Ensure token is defined and handle gracefully if not
// Extract existing guardrails from metadata
let existingGuardrails: string[] = [];
try {
existingGuardrails = token.metadata?.guardrails || [];
} catch (error) {
console.error("Error extracting guardrails:", error);
}
const initialValues = token ? {
...token,
budget_duration: token.budget_duration,
metadata: metadataString
} : { metadata: metadataString };
metadata: metadataString,
guardrails: existingGuardrails
} : { metadata: metadataString, guardrails: [] };
const handleOk = () => {
@ -508,6 +537,33 @@ const ViewKeyTable: React.FC<ViewKeyTableProps> = ({
>
<InputNumber step={1} precision={1} width={200} />
</Form.Item>
<Form.Item
label={
<span>
Guardrails{' '}
<Tooltip title="Setup your first guardrail">
<a
href="https://docs.litellm.ai/docs/proxy/guardrails/quick_start"
target="_blank"
rel="noopener noreferrer"
onClick={(e) => e.stopPropagation()}
>
<InfoCircleOutlined style={{ marginLeft: '4px' }} />
</a>
</Tooltip>
</span>
}
name="guardrails"
className="mt-8"
help="Select existing guardrails or enter new ones"
>
<Select
mode="tags"
style={{ width: '100%' }}
placeholder="Select or enter guardrails"
options={guardrailsList.map(name => ({ value: name, label: name }))}
/>
</Form.Item>
<Form.Item
label="Metadata (ensure this is valid JSON)"
name="metadata"
@ -747,12 +803,23 @@ const ViewKeyTable: React.FC<ViewKeyTableProps> = ({
// Convert metadata back to an object if it exists and is a string
if (formValues.metadata && typeof formValues.metadata === 'string') {
try {
formValues.metadata = JSON.parse(formValues.metadata);
const parsedMetadata = JSON.parse(formValues.metadata);
// Only add guardrails if they are set in form values
formValues.metadata = {
...parsedMetadata,
...(formValues.guardrails?.length > 0 ? { guardrails: formValues.guardrails } : {})
};
} catch (error) {
console.error("Error parsing metadata JSON:", error);
message.error("Invalid metadata JSON for formValue " + formValues.metadata);
return;
}
} else {
// If metadata is not a string (or doesn't exist), only add guardrails if they are set
formValues.metadata = {
...(formValues.metadata || {}),
...(formValues.guardrails?.length > 0 ? { guardrails: formValues.guardrails } : {})
};
}
// Convert the budget_duration back to the API expected format
@ -772,6 +839,7 @@ const ViewKeyTable: React.FC<ViewKeyTableProps> = ({
console.log("handleEditSubmit:", formValues);
try {
let newKeyValues = await keyUpdateCall(accessToken, formValues);
console.log("handleEditSubmit: newKeyValues", newKeyValues);
@ -786,9 +854,14 @@ const ViewKeyTable: React.FC<ViewKeyTableProps> = ({
setEditModalVisible(false);
setSelectedToken(null);
} catch (error) {
console.error("Error updating key:", error);
message.error("Failed to update key");
}
};
const handleDelete = async (token: any) => {
console.log("handleDelete:", token);
if (token.token == null) {