Skip to content

Commit 257a3b4

Browse files
author
Divyansh Agrawal
committed
feat: add budget enforcement for API keys and enhance model allowlists
- Implemented budget enforcement logic in TenantAuthMiddleware to reject requests exceeding daily, weekly, and monthly limits. - Updated APIKeyRepo to include daily and weekly budget limits in the database schema and queries. - Enhanced ProjectRepo and TeamRepo to support allowed models as JSONB fields. - Created a new ModelSelector component for selecting models with search and filtering capabilities. - Added command and popover UI components for improved user interaction. - Updated seed data to reflect changes in model pricing and removed deprecated entries. - Added database migration to introduce new columns for budget limits and model allowlists.
1 parent e0f5553 commit 257a3b4

26 files changed

Lines changed: 1959 additions & 453 deletions

‎dashboard/src/components/ModelSelector.tsx‎

Lines changed: 239 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,239 @@
1+
import { useState, useMemo, useRef } from "react"
2+
import { useQuery } from "@tanstack/react-query"
3+
import { AdminAPI, type ModelPrice } from "@/lib/api"
4+
import {
5+
Check,
6+
ChevronDown,
7+
X,
8+
SearchX
9+
} from "lucide-react"
10+
import { cn } from "@/lib/utils"
11+
import { Button } from "@/components/ui/button"
12+
import { Badge } from "@/components/ui/badge"
13+
import {
14+
Command,
15+
CommandInput,
16+
CommandList,
17+
} from "@/components/ui/command"
18+
import {
19+
Popover,
20+
PopoverContent,
21+
PopoverTrigger,
22+
} from "@/components/ui/popover"
23+
24+
interface ModelSelectorProps {
25+
selectedModels: string[]
26+
onChange: (models: string[]) => void
27+
className?: string
28+
}
29+
30+
export function ModelSelector({ selectedModels, onChange, className }: ModelSelectorProps) {
31+
const [open, setOpen] = useState(false)
32+
const [search, setSearch] = useState("")
33+
const [activeProvider, setActiveProvider] = useState<string | null>(null)
34+
const popoverContainerRef = useRef<HTMLDivElement | null>(null)
35+
36+
const { data: models = [] } = useQuery({
37+
queryKey: ["modelPrices"],
38+
queryFn: AdminAPI.getModelPrices,
39+
})
40+
41+
const uniqueModels = useMemo(() => {
42+
const seen = new Set<string>();
43+
return models.filter(m => {
44+
const key = `${m.provider}:${m.model_name}`;
45+
if (seen.has(key)) return false;
46+
seen.add(key);
47+
return true;
48+
});
49+
}, [models]);
50+
51+
const providers = useMemo(() => {
52+
const set = new Set<string>()
53+
uniqueModels.forEach(m => set.add(m.provider))
54+
return Array.from(set).sort()
55+
}, [uniqueModels])
56+
57+
const filteredModels = useMemo(() => {
58+
let result = uniqueModels
59+
if (search) {
60+
const s = search.toLowerCase()
61+
result = result.filter(m =>
62+
m.model_name.toLowerCase().includes(s) ||
63+
m.provider.toLowerCase().includes(s)
64+
)
65+
}
66+
return result
67+
}, [uniqueModels, search])
68+
69+
const groupedModels = useMemo(() => {
70+
const groups: Record<string, ModelPrice[]> = {}
71+
filteredModels.forEach(m => {
72+
if (!groups[m.provider]) groups[m.provider] = []
73+
groups[m.provider].push(m)
74+
})
75+
return groups
76+
}, [filteredModels])
77+
78+
const toggleModel = (modelName: string) => {
79+
if (selectedModels.includes(modelName)) {
80+
onChange(selectedModels.filter(m => m !== modelName))
81+
} else {
82+
onChange([...selectedModels, modelName])
83+
}
84+
}
85+
86+
const removeModel = (modelName: string) => {
87+
onChange(selectedModels.filter(m => m !== modelName))
88+
}
89+
90+
return (
91+
<div ref={popoverContainerRef} className={cn("space-y-2", className)}>
92+
<Popover open={open} onOpenChange={setOpen}>
93+
<PopoverTrigger asChild>
94+
<div
95+
role="combobox"
96+
aria-expanded={open}
97+
tabIndex={0}
98+
onKeyDown={(e) => {
99+
if (e.key === "Enter" || e.key === " ") {
100+
e.preventDefault()
101+
setOpen((v) => !v)
102+
}
103+
}}
104+
className="relative flex items-center h-10 min-w-0 px-3 pr-12 bg-zinc-100/[0.03] border border-white/10 rounded-[12px] transition-all focus-within:ring-1 focus-within:ring-emerald-500/30 cursor-pointer overflow-hidden"
105+
>
106+
{selectedModels.length === 0 ? (
107+
<span className="text-sm text-zinc-500 truncate">All models allowed</span>
108+
) : (
109+
<div className="flex min-w-0 items-center gap-1.5 overflow-x-auto whitespace-nowrap pr-1">
110+
{selectedModels.map(model => (
111+
<Badge
112+
key={model}
113+
variant="secondary"
114+
className="shrink-0 bg-emerald-500/10 hover:bg-emerald-500/20 text-emerald-300 border-emerald-500/20 py-0.5 pr-1 pl-2 gap-1 rounded-[8px]"
115+
>
116+
<span className="text-[11px] font-medium">{model}</span>
117+
<button
118+
onMouseDown={(e) => e.stopPropagation()}
119+
onClick={(e) => { e.preventDefault(); e.stopPropagation(); removeModel(model); }}
120+
className="hover:text-emerald-100 p-0.5 rounded-full hover:bg-emerald-500/20 transition-colors"
121+
>
122+
<X size={10} />
123+
</button>
124+
</Badge>
125+
))}
126+
</div>
127+
)}
128+
<span
129+
className="absolute right-2 top-1/2 -translate-y-1/2 h-7 w-7 rounded-[8px] border border-white/10 bg-zinc-100/[0.03] text-zinc-400 transition-colors flex items-center justify-center"
130+
aria-hidden
131+
>
132+
<ChevronDown className="h-4 w-4 opacity-80 transition-transform duration-200" style={{ transform: open ? 'rotate(180deg)' : 'none' }} />
133+
</span>
134+
</div>
135+
</PopoverTrigger>
136+
<PopoverContent
137+
container={popoverContainerRef.current}
138+
className="w-[min(420px,calc(100vw-2rem))] p-0 bg-[#0C0C0D] border-white/10 rounded-[16px] shadow-[0_20px_80px_rgba(0,0,0,0.65)] overflow-hidden max-h-[min(520px,var(--radix-popover-content-available-height))] flex flex-col"
139+
side="bottom"
140+
align="start"
141+
sideOffset={8}
142+
collisionPadding={12}
143+
>
144+
<Command shouldFilter={false} className="bg-transparent flex min-h-0 flex-1 flex-col">
145+
<CommandInput
146+
placeholder="Search models..."
147+
value={search}
148+
onValueChange={setSearch}
149+
className="h-11 border-none bg-transparent text-zinc-100 placeholder:text-zinc-500"
150+
/>
151+
<div className="flex min-h-0 flex-1">
152+
{/* Waterfall Sidebar - Only show if not searching */}
153+
{!search && (
154+
<div className="w-[140px] shrink-0 border-r border-white/5 bg-zinc-100/[0.01] overflow-y-auto p-2 space-y-0.5">
155+
<button
156+
onClick={() => setActiveProvider(null)}
157+
className={cn(
158+
"w-full text-left px-3 py-2 rounded-[8px] text-[10px] font-bold uppercase tracking-wider transition-all",
159+
activeProvider === null ? "bg-white text-zinc-950" : "text-zinc-500 hover:text-zinc-300"
160+
)}
161+
>
162+
All
163+
</button>
164+
<div className="px-3 pt-3 pb-1 text-[9px] font-black text-zinc-600 uppercase tracking-[0.2em]">Providers</div>
165+
{providers.map(provider => (
166+
<button
167+
key={provider}
168+
onClick={() => setActiveProvider(provider)}
169+
className={cn(
170+
"w-full text-left px-3 py-2 rounded-[8px] text-xs font-medium transition-all",
171+
activeProvider === provider ? "bg-white/10 text-white" : "text-zinc-500 hover:text-zinc-400 hover:bg-white/[0.02]"
172+
)}
173+
>
174+
{provider}
175+
</button>
176+
))}
177+
</div>
178+
)}
179+
180+
{/* Model List */}
181+
<CommandList className="flex-1 min-w-0 min-h-0 max-h-none overflow-y-auto bg-transparent">
182+
<div className="p-2">
183+
{filteredModels.length === 0 && (
184+
<div className="py-12 flex flex-col items-center justify-center text-zinc-600">
185+
<SearchX size={32} className="mb-2 opacity-20" />
186+
<p className="text-sm">No models found</p>
187+
</div>
188+
)}
189+
{Object.entries(groupedModels)
190+
.filter(([p]) => !activeProvider || p === activeProvider)
191+
.map(([provider, providerModels]) => (
192+
<div key={provider} className="mb-3">
193+
<div className="px-2 pb-1 pt-2 text-[10px] font-black uppercase tracking-[0.15em] text-zinc-600">{provider}</div>
194+
{providerModels.map(model => {
195+
const selected = selectedModels.includes(model.model_name);
196+
return (
197+
<button
198+
key={model.model_name}
199+
onClick={() => toggleModel(model.model_name)}
200+
className={cn(
201+
"w-full flex items-center gap-3 rounded-[8px] px-2 py-2 text-sm transition-colors mb-0.5 text-left",
202+
selected
203+
? "bg-white/10 text-white"
204+
: "text-zinc-400 hover:bg-white/5 hover:text-zinc-100"
205+
)}
206+
>
207+
<div className={cn(
208+
"shrink-0 w-4 h-4 rounded-[4px] border flex items-center justify-center transition-all",
209+
selected ? "border-white text-white" : "border-white/20"
210+
)}>
211+
{selected && <Check size={12} strokeWidth={4} />}
212+
</div>
213+
<span className="font-medium truncate">{model.model_name}</span>
214+
</button>
215+
);
216+
})}
217+
</div>
218+
))}
219+
</div>
220+
</CommandList>
221+
</div>
222+
</Command>
223+
<div className="p-3 bg-white/5 border-t border-white/5 flex items-center justify-between">
224+
<span className="text-[11px] text-zinc-500 font-medium">
225+
{selectedModels.length} selected
226+
</span>
227+
<Button
228+
size="sm"
229+
onClick={() => setOpen(false)}
230+
className="h-8 px-4 rounded-[8px] bg-zinc-100 text-zinc-950 hover:bg-zinc-200 text-[11px] font-bold"
231+
>
232+
Done
233+
</Button>
234+
</div>
235+
</PopoverContent>
236+
</Popover>
237+
</div>
238+
)
239+
}

0 commit comments

Comments
 (0)