mirror of
https://github.com/fosrl/pangolin.git
synced 2026-08-13 16:00:02 +02:00
add virtual api keys to budgets
This commit is contained in:
@@ -1697,6 +1697,8 @@
|
||||
"virtualApiKeysErrorFetchSecret": "Error loading secret",
|
||||
"virtualApiKeysErrorFetchSecretDescription": "Failed to load the virtual API key secret",
|
||||
"virtualApiKeysFilterUnassigned": "Unassigned",
|
||||
"virtualApiKeysInferenceBudget": "Inference Budget",
|
||||
"virtualApiKeysInferenceBudgetDescription": "Configure how this key restricts AI usage based on spending or token limits",
|
||||
"myVirtualApiKeysTitle": "Your API Keys",
|
||||
"myVirtualApiKeysDescription": "View your identity key and any virtual API keys attributed to you in this organization",
|
||||
"myVirtualApiKeysResourceTitle": "Your API Keys for {resourceName}",
|
||||
|
||||
@@ -1776,6 +1776,10 @@ export const aiBudgets = pgTable(
|
||||
roleId: integer("roleId").references(() => roles.roleId, {
|
||||
onDelete: "cascade"
|
||||
}),
|
||||
virtualApiKeyId: varchar("virtualApiKeyId").references(
|
||||
() => virtualApiKeys.virtualApiKeyId,
|
||||
{ onDelete: "cascade" }
|
||||
),
|
||||
amount: real("amount").notNull(),
|
||||
unit: varchar("unit").$type<"usd" | "tokens">().notNull(),
|
||||
period: varchar("period")
|
||||
@@ -1806,7 +1810,12 @@ export const aiBudgets = pgTable(
|
||||
t.unit,
|
||||
t.period
|
||||
),
|
||||
unique("ai_budget_role_uniq").on(t.roleId, t.unit, t.period)
|
||||
unique("ai_budget_role_uniq").on(t.roleId, t.unit, t.period),
|
||||
unique("ai_budget_virtual_api_key_uniq").on(
|
||||
t.virtualApiKeyId,
|
||||
t.unit,
|
||||
t.period
|
||||
)
|
||||
]
|
||||
);
|
||||
|
||||
|
||||
@@ -1762,6 +1762,10 @@ export const aiBudgets = sqliteTable(
|
||||
roleId: integer("roleId").references(() => roles.roleId, {
|
||||
onDelete: "cascade"
|
||||
}),
|
||||
virtualApiKeyId: text("virtualApiKeyId").references(
|
||||
() => virtualApiKeys.virtualApiKeyId,
|
||||
{ onDelete: "cascade" }
|
||||
),
|
||||
amount: real("amount").notNull(),
|
||||
unit: text("unit").$type<"usd" | "tokens">().notNull(),
|
||||
period: text("period")
|
||||
@@ -1794,7 +1798,12 @@ export const aiBudgets = sqliteTable(
|
||||
t.unit,
|
||||
t.period
|
||||
),
|
||||
unique("ai_budget_role_uniq").on(t.roleId, t.unit, t.period)
|
||||
unique("ai_budget_role_uniq").on(t.roleId, t.unit, t.period),
|
||||
unique("ai_budget_virtual_api_key_uniq").on(
|
||||
t.virtualApiKeyId,
|
||||
t.unit,
|
||||
t.period
|
||||
)
|
||||
]
|
||||
);
|
||||
|
||||
|
||||
@@ -53,7 +53,8 @@ function applicableBudgetsCacheKey(ctx: BudgetScopeContext): string {
|
||||
ctx.requestedModel,
|
||||
ctx.resourceId ?? "",
|
||||
ctx.siteResourceId ?? "",
|
||||
roleKey
|
||||
roleKey,
|
||||
ctx.virtualApiKeyId ?? ""
|
||||
].join(":");
|
||||
}
|
||||
|
||||
@@ -83,6 +84,7 @@ export type BudgetScopeContext = {
|
||||
siteResourceId: number | null;
|
||||
roleIds: number[];
|
||||
requestUserId: string | null;
|
||||
virtualApiKeyId: string | null;
|
||||
};
|
||||
|
||||
/**
|
||||
@@ -142,6 +144,11 @@ async function fetchApplicableBudgets(
|
||||
if (ctx.roleIds.length > 0) {
|
||||
scopeConditions.push(inArray(aiBudgets.roleId, ctx.roleIds));
|
||||
}
|
||||
if (ctx.virtualApiKeyId != null) {
|
||||
scopeConditions.push(
|
||||
eq(aiBudgets.virtualApiKeyId, ctx.virtualApiKeyId)
|
||||
);
|
||||
}
|
||||
|
||||
return db
|
||||
.select()
|
||||
@@ -274,6 +281,17 @@ export async function sumUsageForBudget(
|
||||
);
|
||||
}
|
||||
|
||||
if (budget.virtualApiKeyId != null) {
|
||||
return sumUsageAmount(
|
||||
and(
|
||||
eq(aiUsageRecords.orgId, ctx.orgId),
|
||||
eq(aiUsageRecords.virtualApiKeyId, budget.virtualApiKeyId),
|
||||
gte(aiUsageRecords.createdAt, start)
|
||||
)!,
|
||||
budget.unit
|
||||
);
|
||||
}
|
||||
|
||||
return 0;
|
||||
}
|
||||
|
||||
|
||||
@@ -7,7 +7,8 @@ import {
|
||||
db,
|
||||
resources,
|
||||
roles,
|
||||
siteResources
|
||||
siteResources,
|
||||
virtualApiKeys
|
||||
} from "@server/db";
|
||||
import response from "@server/lib/response";
|
||||
import HttpCode from "@server/types/HttpCode";
|
||||
@@ -35,6 +36,7 @@ const bodySchema = z
|
||||
resourceId: z.coerce.number().int().positive().optional(),
|
||||
siteResourceId: z.coerce.number().int().positive().optional(),
|
||||
roleId: z.coerce.number().int().positive().optional(),
|
||||
virtualApiKeyId: z.string().nonempty().optional(),
|
||||
amount: z.number().positive(),
|
||||
unit: aiBudgetUnitSchema,
|
||||
period: aiBudgetPeriodSchema.optional().default("monthly"),
|
||||
@@ -98,6 +100,7 @@ export async function createAiBudget(
|
||||
resourceId,
|
||||
siteResourceId,
|
||||
roleId,
|
||||
virtualApiKeyId,
|
||||
amount,
|
||||
unit,
|
||||
period,
|
||||
@@ -189,6 +192,22 @@ export async function createAiBudget(
|
||||
}
|
||||
}
|
||||
|
||||
if (virtualApiKeyId !== undefined) {
|
||||
const [key] = await db
|
||||
.select({ orgId: virtualApiKeys.orgId })
|
||||
.from(virtualApiKeys)
|
||||
.where(eq(virtualApiKeys.virtualApiKeyId, virtualApiKeyId))
|
||||
.limit(1);
|
||||
if (!key || key.orgId !== orgId) {
|
||||
return next(
|
||||
createHttpError(
|
||||
HttpCode.NOT_FOUND,
|
||||
`Virtual API key with ID ${virtualApiKeyId} not found in this organization`
|
||||
)
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
const scopeCondition =
|
||||
providerId !== undefined
|
||||
? eq(aiBudgets.providerId, providerId)
|
||||
@@ -200,14 +219,17 @@ export async function createAiBudget(
|
||||
? eq(aiBudgets.siteResourceId, siteResourceId)
|
||||
: roleId !== undefined
|
||||
? eq(aiBudgets.roleId, roleId)
|
||||
: and(
|
||||
eq(aiBudgets.orgId, orgId),
|
||||
isNull(aiBudgets.providerId),
|
||||
isNull(aiBudgets.modelId),
|
||||
isNull(aiBudgets.resourceId),
|
||||
isNull(aiBudgets.siteResourceId),
|
||||
isNull(aiBudgets.roleId)
|
||||
);
|
||||
: virtualApiKeyId !== undefined
|
||||
? eq(aiBudgets.virtualApiKeyId, virtualApiKeyId)
|
||||
: and(
|
||||
eq(aiBudgets.orgId, orgId),
|
||||
isNull(aiBudgets.providerId),
|
||||
isNull(aiBudgets.modelId),
|
||||
isNull(aiBudgets.resourceId),
|
||||
isNull(aiBudgets.siteResourceId),
|
||||
isNull(aiBudgets.roleId),
|
||||
isNull(aiBudgets.virtualApiKeyId)
|
||||
);
|
||||
|
||||
const [existing] = await db
|
||||
.select({ budgetId: aiBudgets.budgetId })
|
||||
@@ -239,6 +261,7 @@ export async function createAiBudget(
|
||||
resourceId: resourceId ?? null,
|
||||
siteResourceId: siteResourceId ?? null,
|
||||
roleId: roleId ?? null,
|
||||
virtualApiKeyId: virtualApiKeyId ?? null,
|
||||
amount,
|
||||
unit,
|
||||
period,
|
||||
|
||||
@@ -5,6 +5,7 @@ export * from "./listAiBudgetsForModel";
|
||||
export * from "./listAiBudgetsForResource";
|
||||
export * from "./listAiBudgetsForSiteResource";
|
||||
export * from "./listAiBudgetsForRole";
|
||||
export * from "./listAiBudgetsForVirtualApiKey";
|
||||
export * from "./getAiBudget";
|
||||
export * from "./updateAiBudget";
|
||||
export * from "./deleteAiBudget";
|
||||
|
||||
@@ -0,0 +1,69 @@
|
||||
import { Request, Response, NextFunction } from "express";
|
||||
import { z } from "zod";
|
||||
import { aiBudgets, db } from "@server/db";
|
||||
import response from "@server/lib/response";
|
||||
import HttpCode from "@server/types/HttpCode";
|
||||
import createHttpError from "http-errors";
|
||||
import logger from "@server/logger";
|
||||
import { fromError } from "zod-validation-error";
|
||||
import { OpenAPITags, registry } from "@server/openApi";
|
||||
import { asc, eq } from "drizzle-orm";
|
||||
import type { ListAiBudgetsByScopeResponse } from "@server/routers/aiBudget/types";
|
||||
|
||||
const paramsSchema = z.strictObject({
|
||||
virtualApiKeyId: z.string().nonempty()
|
||||
});
|
||||
|
||||
registry.registerPath({
|
||||
method: "get",
|
||||
path: "/virtual-api-key/{virtualApiKeyId}/ai-budgets",
|
||||
description: "List AI budgets scoped to a virtual API key.",
|
||||
tags: [OpenAPITags.AiBudget],
|
||||
request: {
|
||||
params: paramsSchema
|
||||
},
|
||||
responses: {
|
||||
200: {
|
||||
description: "Successful response"
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
export async function listAiBudgetsForVirtualApiKey(
|
||||
req: Request,
|
||||
res: Response,
|
||||
next: NextFunction
|
||||
): Promise<any> {
|
||||
try {
|
||||
const parsedParams = paramsSchema.safeParse(req.params);
|
||||
if (!parsedParams.success) {
|
||||
return next(
|
||||
createHttpError(
|
||||
HttpCode.BAD_REQUEST,
|
||||
fromError(parsedParams.error).toString()
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
const { virtualApiKeyId } = parsedParams.data;
|
||||
|
||||
const budgets = await db
|
||||
.select()
|
||||
.from(aiBudgets)
|
||||
.where(eq(aiBudgets.virtualApiKeyId, virtualApiKeyId))
|
||||
.orderBy(asc(aiBudgets.budgetId));
|
||||
|
||||
return response<ListAiBudgetsByScopeResponse>(res, {
|
||||
data: { budgets },
|
||||
success: true,
|
||||
error: false,
|
||||
message: "AI budgets retrieved successfully",
|
||||
status: HttpCode.OK
|
||||
});
|
||||
} catch (error) {
|
||||
logger.error(error);
|
||||
return next(
|
||||
createHttpError(HttpCode.INTERNAL_SERVER_ERROR, "An error occurred")
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -7,7 +7,8 @@ import {
|
||||
db,
|
||||
resources,
|
||||
roles,
|
||||
siteResources
|
||||
siteResources,
|
||||
virtualApiKeys
|
||||
} from "@server/db";
|
||||
import response from "@server/lib/response";
|
||||
import HttpCode from "@server/types/HttpCode";
|
||||
@@ -34,6 +35,7 @@ const bodySchema = z.strictObject({
|
||||
resourceId: z.coerce.number().int().positive().nullable().optional(),
|
||||
siteResourceId: z.coerce.number().int().positive().nullable().optional(),
|
||||
roleId: z.coerce.number().int().positive().nullable().optional(),
|
||||
virtualApiKeyId: z.string().nonempty().nullable().optional(),
|
||||
amount: z.number().positive().optional(),
|
||||
unit: aiBudgetUnitSchema.optional(),
|
||||
period: aiBudgetPeriodSchema.optional(),
|
||||
@@ -128,6 +130,10 @@ export async function updateAiBudget(
|
||||
: existing.siteResourceId;
|
||||
const nextRoleId =
|
||||
body.roleId !== undefined ? body.roleId : existing.roleId;
|
||||
const nextVirtualApiKeyId =
|
||||
body.virtualApiKeyId !== undefined
|
||||
? body.virtualApiKeyId
|
||||
: existing.virtualApiKeyId;
|
||||
const nextUnit = body.unit !== undefined ? body.unit : existing.unit;
|
||||
const nextPeriod =
|
||||
body.period !== undefined ? body.period : existing.period;
|
||||
@@ -138,7 +144,8 @@ export async function updateAiBudget(
|
||||
modelId: z.number().nullable().optional(),
|
||||
resourceId: z.number().nullable().optional(),
|
||||
siteResourceId: z.number().nullable().optional(),
|
||||
roleId: z.number().nullable().optional()
|
||||
roleId: z.number().nullable().optional(),
|
||||
virtualApiKeyId: z.string().nullable().optional()
|
||||
})
|
||||
.superRefine((data, ctx) => refineBudgetScopeFields(data, ctx))
|
||||
.safeParse({
|
||||
@@ -146,7 +153,8 @@ export async function updateAiBudget(
|
||||
modelId: nextModelId,
|
||||
resourceId: nextResourceId,
|
||||
siteResourceId: nextSiteResourceId,
|
||||
roleId: nextRoleId
|
||||
roleId: nextRoleId,
|
||||
virtualApiKeyId: nextVirtualApiKeyId
|
||||
});
|
||||
|
||||
if (!scopeValidation.success) {
|
||||
@@ -245,6 +253,30 @@ export async function updateAiBudget(
|
||||
}
|
||||
}
|
||||
|
||||
if (
|
||||
body.virtualApiKeyId !== undefined &&
|
||||
body.virtualApiKeyId !== null
|
||||
) {
|
||||
const [key] = await db
|
||||
.select({ orgId: virtualApiKeys.orgId })
|
||||
.from(virtualApiKeys)
|
||||
.where(
|
||||
eq(
|
||||
virtualApiKeys.virtualApiKeyId,
|
||||
body.virtualApiKeyId
|
||||
)
|
||||
)
|
||||
.limit(1);
|
||||
if (!key || key.orgId !== orgId) {
|
||||
return next(
|
||||
createHttpError(
|
||||
HttpCode.NOT_FOUND,
|
||||
`Virtual API key with ID ${body.virtualApiKeyId} not found in this organization`
|
||||
)
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
const scopeCondition =
|
||||
nextProviderId !== null
|
||||
? eq(aiBudgets.providerId, nextProviderId)
|
||||
@@ -256,14 +288,20 @@ export async function updateAiBudget(
|
||||
? eq(aiBudgets.siteResourceId, nextSiteResourceId)
|
||||
: nextRoleId !== null
|
||||
? eq(aiBudgets.roleId, nextRoleId)
|
||||
: and(
|
||||
eq(aiBudgets.orgId, orgId),
|
||||
isNull(aiBudgets.providerId),
|
||||
isNull(aiBudgets.modelId),
|
||||
isNull(aiBudgets.resourceId),
|
||||
isNull(aiBudgets.siteResourceId),
|
||||
isNull(aiBudgets.roleId)
|
||||
);
|
||||
: nextVirtualApiKeyId !== null
|
||||
? eq(
|
||||
aiBudgets.virtualApiKeyId,
|
||||
nextVirtualApiKeyId
|
||||
)
|
||||
: and(
|
||||
eq(aiBudgets.orgId, orgId),
|
||||
isNull(aiBudgets.providerId),
|
||||
isNull(aiBudgets.modelId),
|
||||
isNull(aiBudgets.resourceId),
|
||||
isNull(aiBudgets.siteResourceId),
|
||||
isNull(aiBudgets.roleId),
|
||||
isNull(aiBudgets.virtualApiKeyId)
|
||||
);
|
||||
|
||||
const [conflict] = await db
|
||||
.select({ budgetId: aiBudgets.budgetId })
|
||||
@@ -305,6 +343,9 @@ export async function updateAiBudget(
|
||||
if (body.roleId !== undefined) {
|
||||
updateData.roleId = body.roleId;
|
||||
}
|
||||
if (body.virtualApiKeyId !== undefined) {
|
||||
updateData.virtualApiKeyId = body.virtualApiKeyId;
|
||||
}
|
||||
if (body.amount !== undefined) {
|
||||
updateData.amount = body.amount;
|
||||
}
|
||||
|
||||
@@ -20,6 +20,7 @@ export function refineBudgetScopeFields(
|
||||
resourceId?: number | null;
|
||||
siteResourceId?: number | null;
|
||||
roleId?: number | null;
|
||||
virtualApiKeyId?: string | null;
|
||||
},
|
||||
ctx: z.RefinementCtx
|
||||
) {
|
||||
@@ -28,7 +29,8 @@ export function refineBudgetScopeFields(
|
||||
data.modelId,
|
||||
data.resourceId,
|
||||
data.siteResourceId,
|
||||
data.roleId
|
||||
data.roleId,
|
||||
data.virtualApiKeyId
|
||||
];
|
||||
|
||||
const setCount = scopeFields.filter(
|
||||
@@ -39,7 +41,7 @@ export function refineBudgetScopeFields(
|
||||
ctx.addIssue({
|
||||
code: "custom",
|
||||
message:
|
||||
"Only one of providerId, modelId, resourceId, siteResourceId, or roleId may be set on a budget",
|
||||
"Only one of providerId, modelId, resourceId, siteResourceId, roleId, or virtualApiKeyId may be set on a budget",
|
||||
path: ["providerId"]
|
||||
});
|
||||
}
|
||||
|
||||
@@ -830,7 +830,8 @@ export async function handleAiGatewayProxy(
|
||||
resourceId,
|
||||
siteResourceId,
|
||||
roleIds: requestUser?.roleIds ?? [],
|
||||
requestUserId: requestUser?.userId ?? null
|
||||
requestUserId: requestUser?.userId ?? null,
|
||||
virtualApiKeyId: identity.virtualApiKeyId
|
||||
});
|
||||
appliedBudgets = budgetCheck.budgets;
|
||||
|
||||
|
||||
@@ -1791,6 +1791,13 @@ authenticated.get(
|
||||
aiBudget.listAiBudgetsForRole
|
||||
);
|
||||
|
||||
authenticated.get(
|
||||
"/virtual-api-key/:virtualApiKeyId/ai-budgets",
|
||||
verifyVirtualApiKeyAccess,
|
||||
verifyUserHasAction(ActionsEnum.listAiBudgets),
|
||||
aiBudget.listAiBudgetsForVirtualApiKey
|
||||
);
|
||||
|
||||
authenticated.get(
|
||||
"/org/:orgId/labels",
|
||||
verifyOrgAccess,
|
||||
|
||||
@@ -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>
|
||||
)}
|
||||
|
||||
@@ -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>
|
||||
|
||||
@@ -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"
|
||||
}
|
||||
};
|
||||
|
||||
|
||||
Reference in New Issue
Block a user