Skip to content

Commit 83654c5

Browse files
mubashir1osmanimubashir1osmani
committed
feat: add UI support for displaying and managing guardrails on virtual keys
1 parent 6cd5afa commit 83654c5

File tree

8 files changed

+364
-0
lines changed

8 files changed

+364
-0
lines changed

‎enterprise/litellm_enterprise/proxy/guardrails/endpoints.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -97,7 +97,9 @@ async def get_guardrails_for_virtual_key(
9797
"""
9898
Get all guardrails associated with a virtual key
9999
"""
100+
verbose_proxy_logger.debug(f"Getting guardrails for virtual key: {virtual_key_id}")
100101
guardrails = IN_MEMORY_GUARDRAIL_HANDLER.get_guardrails_for_virtual_key(virtual_key_id)
102+
verbose_proxy_logger.debug(f"Found {len(guardrails)} guardrails for virtual key {virtual_key_id}")
101103

102104
return VirtualKeyGuardrailsResponse(
103105
virtual_key_id=virtual_key_id,

‎litellm/proxy/guardrails/guardrail_registry.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -567,6 +567,7 @@ def get_guardrails_for_virtual_key(self, virtual_key_id: str) -> List[Guardrail]
567567
Get all guardrails associated with a virtual key
568568
"""
569569
guardrail_ids = self.virtual_key_to_guardrails.get(virtual_key_id, [])
570+
verbose_proxy_logger.debug(f"Getting guardrails for virtual key {virtual_key_id}: {guardrail_ids}")
570571
return [self.IN_MEMORY_GUARDRAILS[gid] for gid in guardrail_ids if gid in self.IN_MEMORY_GUARDRAILS]
571572

572573

‎ui/litellm-dashboard/src/components/virtual_keys/virtual_key_detail.tsx‎

Lines changed: 160 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,160 @@
1+
import React, { useState, useEffect } from 'react';
2+
import { Card, Flex, Text, Heading, Button, Dialog, TextField, Select, Box, Tabs } from '@radix-ui/themes';
3+
import { toast } from 'react-hot-toast';
4+
import { VirtualKey } from '../../types/virtual_key';
5+
import { fetchVirtualKey, updateVirtualKey } from '../../services/virtual_keys';
6+
import { VirtualKeyGuardrails } from './virtual_key_guardrails';
7+
import { fetchGuardrails } from '../../services/guardrails';
8+
import { Guardrail } from '../../types/guardrail';
9+
import { associateGuardrailWithVirtualKey } from '../../services/virtual_key_guardrails';
10+
11+
interface VirtualKeyDetailProps {
12+
virtualKeyId: string;
13+
onClose: () => void;
14+
}
15+
16+
export const VirtualKeyDetail: React.FC<VirtualKeyDetailProps> = ({ virtualKeyId, onClose }) => {
17+
const [virtualKey, setVirtualKey] = useState<VirtualKey | null>(null);
18+
const [loading, setLoading] = useState(true);
19+
const [availableGuardrails, setAvailableGuardrails] = useState<Guardrail[]>([]);
20+
const [selectedGuardrailId, setSelectedGuardrailId] = useState<string>('');
21+
const [refreshTrigger, setRefreshTrigger] = useState(0);
22+
23+
useEffect(() => {
24+
const loadVirtualKey = async () => {
25+
try {
26+
const data = await fetchVirtualKey(virtualKeyId);
27+
setVirtualKey(data);
28+
} catch (error) {
29+
console.error('Error loading virtual key:', error);
30+
toast.error('Failed to load virtual key details');
31+
} finally {
32+
setLoading(false);
33+
}
34+
};
35+
36+
const loadGuardrails = async () => {
37+
try {
38+
const data = await fetchGuardrails();
39+
setAvailableGuardrails(data.guardrails);
40+
} catch (error) {
41+
console.error('Error loading guardrails:', error);
42+
}
43+
};
44+
45+
loadVirtualKey();
46+
loadGuardrails();
47+
}, [virtualKeyId]);
48+
49+
const handleAddGuardrail = async () => {
50+
if (!selectedGuardrailId) {
51+
toast.error('Please select a guardrail');
52+
return;
53+
}
54+
55+
try {
56+
await associateGuardrailWithVirtualKey(virtualKeyId, selectedGuardrailId);
57+
toast.success('Guardrail added to virtual key');
58+
setSelectedGuardrailId('');
59+
setRefreshTrigger(prev => prev + 1);
60+
} catch (error) {
61+
console.error('Error adding guardrail:', error);
62+
toast.error('Failed to add guardrail');
63+
}
64+
};
65+
66+
if (loading) {
67+
return (
68+
<Dialog.Root open={true} onOpenChange={onClose}>
69+
<Dialog.Content>
70+
<Dialog.Title>Virtual Key Details</Dialog.Title>
71+
<Text>Loading...</Text>
72+
</Dialog.Content>
73+
</Dialog.Root>
74+
);
75+
}
76+
77+
if (!virtualKey) {
78+
return (
79+
<Dialog.Root open={true} onOpenChange={onClose}>
80+
<Dialog.Content>
81+
<Dialog.Title>Error</Dialog.Title>
82+
<Text>Virtual key not found</Text>
83+
<Dialog.Close>
84+
<Button>Close</Button>
85+
</Dialog.Close>
86+
</Dialog.Content>
87+
</Dialog.Root>
88+
);
89+
}
90+
91+
return (
92+
<Dialog.Root open={true} onOpenChange={onClose}>
93+
<Dialog.Content style={{ maxWidth: '600px' }}>
94+
<Dialog.Title>Virtual Key: {virtualKey.key_name || 'Unnamed Key'}</Dialog.Title>
95+
96+
<Tabs.Root defaultValue="details">
97+
<Tabs.List>
98+
<Tabs.Trigger value="details">Details</Tabs.Trigger>
99+
<Tabs.Trigger value="guardrails">Guardrails</Tabs.Trigger>
100+
</Tabs.List>
101+
102+
<Box py="4">
103+
<Tabs.Content value="details">
104+
<Card>
105+
<Flex direction="column" gap="3">
106+
<Flex justify="between">
107+
<Text weight="bold">Key ID:</Text>
108+
<Text>{virtualKey.key_id}</Text>
109+
</Flex>
110+
<Flex justify="between">
111+
<Text weight="bold">Team ID:</Text>
112+
<Text>{virtualKey.team_id || 'N/A'}</Text>
113+
</Flex>
114+
<Flex justify="between">
115+
<Text weight="bold">Models:</Text>
116+
<Text>{virtualKey.models?.join(', ') || 'All models'}</Text>
117+
</Flex>
118+
<Flex justify="between">
119+
<Text weight="bold">Spend:</Text>
120+
<Text>${virtualKey.spend?.toFixed(2) || '0.00'}</Text>
121+
</Flex>
122+
</Flex>
123+
</Card>
124+
</Tabs.Content>
125+
126+
<Tabs.Content value="guardrails">
127+
<Flex direction="column" gap="3">
128+
<Card>
129+
<Heading size="3" mb="2">Add Guardrail</Heading>
130+
<Flex gap="2">
131+
<Select.Root value={selectedGuardrailId} onValueChange={setSelectedGuardrailId}>
132+
<Select.Trigger placeholder="Select a guardrail" />
133+
<Select.Content>
134+
{availableGuardrails.map(guardrail => (
135+
<Select.Item key={guardrail.guardrail_id} value={guardrail.guardrail_id}>
136+
{guardrail.guardrail_name || `Guardrail ${guardrail.guardrail_id.substring(0, 8)}`}
137+
</Select.Item>
138+
))}
139+
</Select.Content>
140+
</Select.Root>
141+
<Button onClick={handleAddGuardrail}>Add</Button>
142+
</Flex>
143+
</Card>
144+
145+
<VirtualKeyGuardrails
146+
virtualKeyId={virtualKeyId}
147+
refreshTrigger={refreshTrigger}
148+
/>
149+
</Flex>
150+
</Tabs.Content>
151+
</Box>
152+
</Tabs.Root>
153+
154+
<Dialog.Close asChild>
155+
<Button color="gray" variant="soft" mt="4">Close</Button>
156+
</Dialog.Close>
157+
</Dialog.Content>
158+
</Dialog.Root>
159+
);
160+
};

‎ui/litellm-dashboard/src/components/virtual_keys/virtual_key_guardrails.tsx‎

Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,93 @@
1+
import React, { useEffect, useState } from 'react';
2+
import { Badge, Button, Card, Flex, Heading, Text, Box, Spinner } from '@radix-ui/themes';
3+
import { fetchGuardrailsForVirtualKey, associateGuardrailWithVirtualKey, disassociateGuardrailFromVirtualKey } from '../../services/virtual_key_guardrails';
4+
import { Guardrail } from '../../types/guardrail';
5+
import { toast } from 'react-hot-toast';
6+
7+
interface VirtualKeyGuardrailsProps {
8+
virtualKeyId: string;
9+
refreshTrigger?: number;
10+
}
11+
12+
export const VirtualKeyGuardrails: React.FC<VirtualKeyGuardrailsProps> = ({
13+
virtualKeyId,
14+
refreshTrigger = 0
15+
}) => {
16+
const [guardrails, setGuardrails] = useState<Guardrail[]>([]);
17+
const [loading, setLoading] = useState<boolean>(true);
18+
const [error, setError] = useState<string | null>(null);
19+
20+
useEffect(() => {
21+
const loadGuardrails = async () => {
22+
if (!virtualKeyId) return;
23+
24+
setLoading(true);
25+
try {
26+
const data = await fetchGuardrailsForVirtualKey(virtualKeyId);
27+
setGuardrails(data.guardrails);
28+
setError(null);
29+
} catch (err) {
30+
console.error('Error loading guardrails for virtual key:', err);
31+
setError('Failed to load guardrails');
32+
setGuardrails([]);
33+
} finally {
34+
setLoading(false);
35+
}
36+
};
37+
38+
loadGuardrails();
39+
}, [virtualKeyId, refreshTrigger]);
40+
41+
const handleRemoveGuardrail = async (guardrailId: string) => {
42+
try {
43+
await disassociateGuardrailFromVirtualKey(virtualKeyId, guardrailId);
44+
setGuardrails(guardrails.filter(g => g.guardrail_id !== guardrailId));
45+
toast.success('Guardrail removed from virtual key');
46+
} catch (err) {
47+
console.error('Error removing guardrail:', err);
48+
toast.error('Failed to remove guardrail');
49+
}
50+
};
51+
52+
if (loading) {
53+
return (
54+
<Box py="4">
55+
<Flex align="center" justify="center">
56+
<Spinner size="large" />
57+
</Flex>
58+
</Box>
59+
);
60+
}
61+
62+
if (error) {
63+
return (
64+
<Box py="4">
65+
<Text color="red">{error}</Text>
66+
</Box>
67+
);
68+
}
69+
70+
return (
71+
<Card>
72+
<Heading size="3" mb="2">Attached Guardrails</Heading>
73+
{guardrails.length === 0 ? (
74+
<Text color="gray">No guardrails attached to this virtual key</Text>
75+
) : (
76+
<Flex direction="column" gap="2">
77+
{guardrails.map((guardrail) => (
78+
<Flex key={guardrail.guardrail_id} justify="between" align="center" p="2" style={{ borderBottom: '1px solid var(--gray-4)' }}>
79+
<Flex direction="column" gap="1">
80+
<Text weight="bold">{guardrail.guardrail_name || 'Unnamed Guardrail'}</Text>
81+
<Text size="1" color="gray">Type: {guardrail.litellm_params?.guardrail}</Text>
82+
<Text size="1" color="gray">Mode: {guardrail.litellm_params?.mode}</Text>
83+
</Flex>
84+
<Button color="red" variant="soft" onClick={() => handleRemoveGuardrail(guardrail.guardrail_id)}>
85+
Remove
86+
</Button>
87+
</Flex>
88+
))}
89+
</Flex>
90+
)}
91+
</Card>
92+
);
93+
};

‎ui/litellm-dashboard/src/services/guardrails.ts‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
import { Guardrail } from '../types/guardrail';
2+
3+
interface GuardrailsResponse {
4+
guardrails: Guardrail[];
5+
}
6+
7+
export async function fetchGuardrails(): Promise<GuardrailsResponse> {
8+
const response = await fetch('/api/guardrails');
9+
10+
if (!response.ok) {
11+
const errorText = await response.text();
12+
throw new Error(`Failed to fetch guardrails: ${errorText}`);
13+
}
14+
15+
return await response.json();
16+
}
17+
18+
export async function fetchGuardrail(guardrailId: string): Promise<Guardrail> {
19+
const response = await fetch(`/api/guardrails/${guardrailId}`);
20+
21+
if (!response.ok) {
22+
const errorText = await response.text();
23+
throw new Error(`Failed to fetch guardrail: ${errorText}`);
24+
}
25+
26+
return await response.json();
27+
}

‎ui/litellm-dashboard/src/services/virtual_key_guardrails.ts‎

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,57 @@
1+
import { Guardrail } from '../types/guardrail';
2+
3+
interface VirtualKeyGuardrailsResponse {
4+
virtual_key_id: string;
5+
guardrails: Guardrail[];
6+
}
7+
8+
export async function fetchGuardrailsForVirtualKey(virtualKeyId: string): Promise<VirtualKeyGuardrailsResponse> {
9+
const response = await fetch(`/api/guardrails/virtual_key/${virtualKeyId}`);
10+
11+
if (!response.ok) {
12+
const errorText = await response.text();
13+
throw new Error(`Failed to fetch guardrails: ${errorText}`);
14+
}
15+
16+
return await response.json();
17+
}
18+
19+
export async function associateGuardrailWithVirtualKey(virtualKeyId: string, guardrailId: string): Promise<{ message: string }> {
20+
const response = await fetch('/api/guardrails/virtual_key/associate', {
21+
method: 'POST',
22+
headers: {
23+
'Content-Type': 'application/json',
24+
},
25+
body: JSON.stringify({
26+
virtual_key_id: virtualKeyId,
27+
guardrail_id: guardrailId,
28+
}),
29+
});
30+
31+
if (!response.ok) {
32+
const errorText = await response.text();
33+
throw new Error(`Failed to associate guardrail: ${errorText}`);
34+
}
35+
36+
return await response.json();
37+
}
38+
39+
export async function disassociateGuardrailFromVirtualKey(virtualKeyId: string, guardrailId: string): Promise<{ message: string }> {
40+
const response = await fetch('/api/guardrails/virtual_key/disassociate', {
41+
method: 'POST',
42+
headers: {
43+
'Content-Type': 'application/json',
44+
},
45+
body: JSON.stringify({
46+
virtual_key_id: virtualKeyId,
47+
guardrail_id: guardrailId,
48+
}),
49+
});
50+
51+
if (!response.ok) {
52+
const errorText = await response.text();
53+
throw new Error(`Failed to disassociate guardrail: ${errorText}`);
54+
}
55+
56+
return await response.json();
57+
}

‎ui/litellm-dashboard/src/types/guardrail.ts‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
export interface Guardrail {
2+
guardrail_id: string;
3+
guardrail_name: string | null;
4+
litellm_params: {
5+
guardrail: string;
6+
mode: string;
7+
default_on: boolean;
8+
[key: string]: any;
9+
};
10+
guardrail_info?: Record<string, any>;
11+
}

‎ui/litellm-dashboard/src/types/virtual_key.ts‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
export interface VirtualKey {
2+
key_id: string;
3+
key_name?: string;
4+
team_id?: string;
5+
models?: string[];
6+
spend?: number;
7+
max_budget?: number;
8+
max_parallel_requests?: number;
9+
metadata?: Record<string, any>;
10+
expires?: string;
11+
created_at?: string;
12+
updated_at?: string;
13+
}

0 commit comments

Comments
 (0)