calvin.erfmann/pangolin
1
0
Fork 0
mirror of https://github.com/fosrl/pangolin.git synced 2026-08-14 16:30:15 +02:00
Code Issues Packages Projects Releases Wiki Activity

add virtual api keys to budgets

Browse Source
This commit is contained in:
Owen
2026-08-12 13:21:46 -04:00
parent c33fa8782b
commit 47ae017f94
14 changed files with 415 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/components/CreateVirtualApiKeyForm.tsx
+91 -1
View File
@@ -48,6 +48,12 @@ import {
formatMultiResourcesSelectorLabel
} from "@app/components/multi-resource-selector";
import type { SelectedResource } from "@app/components/resource-selector";
import { HorizontalTabs } from "@app/components/HorizontalTabs";
import {
BudgetRowsFields,
getBudgetRowsErrors,
type BudgetRow
} from "@app/components/BudgetsEditor";
export type CreatedVirtualApiKey = {
virtualApiKeyId: string;
@@ -93,6 +99,10 @@ export default function CreateVirtualApiKeyForm({
const [selectedResources, setSelectedResources] = useState<
SelectedResource[]
>([]);
const [pendingBudgetRows, setPendingBudgetRows] = useState<BudgetRow[]>(
[]
);
const [attemptedBudgetsSave, setAttemptedBudgetsSave] = useState(false);
const formSchema = z.object({
name: z.string().min(1),
@@ -113,9 +123,29 @@ export default function CreateVirtualApiKeyForm({
setAllResources(false);
setSelectedUser(null);
setSelectedResources([]);
setPendingBudgetRows([]);
setAttemptedBudgetsSave(false);
form.reset();
}
function handleFormSubmit(values: z.infer<typeof formSchema>) {
const { conflictingKeys, invalidAmountKeys } =
getBudgetRowsErrors(pendingBudgetRows);
if (conflictingKeys.size > 0 || invalidAmountKeys.size > 0) {
setAttemptedBudgetsSave(true);
toast({
variant: "destructive",
title: t("aiBudgetErrorSave"),
description: conflictingKeys.size
? t("aiBudgetConflictError")
: t("aiBudgetInvalidAmountError")
});
return;
}
return onSubmit(values);
}
async function onSubmit(values: z.infer<typeof formSchema>) {
setLoading(true);
@@ -155,6 +185,33 @@ export default function CreateVirtualApiKeyForm({
);
}
const pendingBudgets = pendingBudgetRows.filter(
(budget) => budget.amount.trim() !== ""
);
if (pendingBudgets.length > 0) {
try {
await Promise.all(
pendingBudgets.map((budget) =>
api.put(`/org/${org.org.orgId}/ai-budget`, {
virtualApiKeyId: key.virtualApiKeyId,
amount: Number(budget.amount),
unit: budget.unit,
period: budget.period
})
)
);
} catch (e) {
toast({
variant: "destructive",
title: t("aiBudgetErrorSave"),
description: formatAxiosError(
e,
t("aiBudgetErrorSave")
)
});
}
}
const resourceLookup = new Map(
selectedResources.map((r) => [
r.resourceId,
@@ -219,10 +276,26 @@ export default function CreateVirtualApiKeyForm({
{!credential && (
<Form {...form}>
<form
onSubmit={form.handleSubmit(onSubmit)}
onSubmit={form.handleSubmit(
handleFormSubmit
)}
className="space-y-4"
id="virtual-api-key-form"
>
<HorizontalTabs
clientSide={true}
defaultTab={0}
items={[
{ title: t("general"), href: "#" },
{
title: t(
"virtualApiKeysInferenceBudget"
),
href: "#"
}
]}
>
<div className="space-y-4 mt-4">
<FormField
control={form.control}
name="name"
@@ -392,6 +465,23 @@ export default function CreateVirtualApiKeyForm({
</div>
)}
</div>
</div>
<div className="space-y-4 mt-4">
<p className="text-sm text-muted-foreground">
{t(
"virtualApiKeysInferenceBudgetDescription"
)}
</p>
<BudgetRowsFields
rows={pendingBudgetRows}
onChange={setPendingBudgetRows}
attemptedSave={
attemptedBudgetsSave
}
/>
</div>
</HorizontalTabs>
</form>
</Form>
)}
src/components/EditVirtualApiKeyForm.tsx
+103 -1
View File
@@ -50,6 +50,16 @@ import { getUserDisplayName } from "@app/lib/getUserDisplayName";
import CopyTextBox from "@app/components/CopyTextBox";
import type { CreatedVirtualApiKey } from "@app/components/CreateVirtualApiKeyForm";
import type { GetVirtualApiKeyResponse } from "@server/routers/virtualApiKey/types";
import { HorizontalTabs } from "@app/components/HorizontalTabs";
import {
BudgetRowsFields,
getBudgetRowsErrors,
rowsFromBudgets,
saveBudgetRows,
type BudgetRow
} from "@app/components/BudgetsEditor";
import { aiBudgetQueries } from "@app/lib/queries";
import { useQuery, useQueryClient } from "@tanstack/react-query";
type FormProps = {
open: boolean;
@@ -93,6 +103,7 @@ export default function EditVirtualApiKeyForm({
const { env } = useEnvContext();
const api = createApiClient({ env });
const t = useTranslations();
const queryClient = useQueryClient();
const [loading, setLoading] = useState(false);
const [selectedUser, setSelectedUser] = useState<SelectedUser | null>(null);
@@ -101,6 +112,19 @@ export default function EditVirtualApiKeyForm({
>([]);
const [credential, setCredential] = useState<string | null>(null);
const [credentialLoading, setCredentialLoading] = useState(false);
const [pendingBudgetRows, setPendingBudgetRows] = useState<BudgetRow[]>(
[]
);
const [attemptedBudgetsSave, setAttemptedBudgetsSave] = useState(false);
const budgetScope = {
type: "virtualApiKey" as const,
id: virtualApiKey?.virtualApiKeyId ?? ""
};
const budgetsQuery = useQuery({
...aiBudgetQueries.scoped({ scope: budgetScope }),
enabled: open && !!virtualApiKey
});
const formSchema = z
.object({
@@ -191,6 +215,32 @@ export default function EditVirtualApiKeyForm({
};
}, [open, virtualApiKey, form]);
useEffect(() => {
if (!open || !budgetsQuery.data) {
return;
}
setPendingBudgetRows(rowsFromBudgets(budgetsQuery.data));
setAttemptedBudgetsSave(false);
}, [open, budgetsQuery.data]);
function handleFormSubmit(values: z.infer<typeof formSchema>) {
const { conflictingKeys, invalidAmountKeys } =
getBudgetRowsErrors(pendingBudgetRows);
if (conflictingKeys.size > 0 || invalidAmountKeys.size > 0) {
setAttemptedBudgetsSave(true);
toast({
variant: "destructive",
title: t("aiBudgetErrorSave"),
description: conflictingKeys.size
? t("aiBudgetConflictError")
: t("aiBudgetInvalidAmountError")
});
return;
}
return onSubmit(values);
}
async function onSubmit(values: z.infer<typeof formSchema>) {
if (!virtualApiKey) {
return;
@@ -223,6 +273,26 @@ export default function EditVirtualApiKeyForm({
if (res?.data.data.virtualApiKey) {
const key = res.data.data.virtualApiKey;
try {
await saveBudgetRows({
api,
orgId: virtualApiKey.orgId,
scope: budgetScope,
existingBudgets: budgetsQuery.data ?? [],
rows: pendingBudgetRows
});
await queryClient.invalidateQueries(
aiBudgetQueries.scoped({ scope: budgetScope })
);
} catch (e) {
toast({
variant: "destructive",
title: t("aiBudgetErrorSave"),
description: formatAxiosError(e, t("aiBudgetErrorSave"))
});
}
const resourceLookup = new Map(
selectedResources.map((r) => [
r.resourceId,
@@ -280,10 +350,26 @@ export default function EditVirtualApiKeyForm({
<div className="flex flex-col gap-y-4 px-1">
<Form {...form}>
<form
onSubmit={form.handleSubmit(onSubmit)}
onSubmit={form.handleSubmit(
handleFormSubmit
)}
className="space-y-4"
id="edit-virtual-api-key-form"
>
<HorizontalTabs
clientSide={true}
defaultTab={0}
items={[
{ title: t("general"), href: "#" },
{
title: t(
"virtualApiKeysInferenceBudget"
),
href: "#"
}
]}
>
<div className="space-y-4 mt-4">
<div className="space-y-2">
<Label>
{t(
@@ -430,6 +516,22 @@ export default function EditVirtualApiKeyForm({
</div>
)}
</div>
</div>
<div className="space-y-4 mt-4">
<p className="text-sm text-muted-foreground">
{t(
"virtualApiKeysInferenceBudgetDescription"
)}
</p>
<BudgetRowsFields
rows={pendingBudgetRows}
onChange={setPendingBudgetRows}
disabled={budgetsQuery.isLoading}
attemptedSave={attemptedBudgetsSave}
/>
</div>
</HorizontalTabs>
</form>
</Form>
</div>
src/lib/aiBudgetScope.ts
+13 -4
View File
@@ -5,11 +5,12 @@ export type AiBudgetScopeType =
| "model"
| "resource"
| "siteResource"
| "role";
| "role"
| "virtualApiKey";
export type AiBudgetScope = {
type: AiBudgetScopeType;
id: number;
id: number | string;
};
export type AiBudgetScopeBodyField =
@@ -17,11 +18,15 @@ export type AiBudgetScopeBodyField =
| "modelId"
| "resourceId"
| "siteResourceId"
| "roleId";
| "roleId"
| "virtualApiKeyId";
const scopeConfig: Record<
AiBudgetScopeType,
{ listPath: (id: number) => string; bodyField: AiBudgetScopeBodyField }
{
listPath: (id: number | string) => string;
bodyField: AiBudgetScopeBodyField;
}
> = {
provider: {
listPath: (id) => `/ai-provider/${id}/ai-budgets`,
@@ -42,6 +47,10 @@ const scopeConfig: Record<
role: {
listPath: (id) => `/role/${id}/ai-budgets`,
bodyField: "roleId"
},
virtualApiKey: {
listPath: (id) => `/virtual-api-key/${id}/ai-budgets`,
bodyField: "virtualApiKeyId"
}
};