calvin.erfmann/pangolin
1
0
Fork 0
mirror of https://github.com/fosrl/pangolin.git synced 2026-08-13 16:00:02 +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
server/routers/aiBudget/createAiBudget.ts
+32 -9
View File
@@ -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,
server/routers/aiBudget/index.ts
+1
View File
@@ -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";
server/routers/aiBudget/listAiBudgetsForVirtualApiKey.ts
+69
View File
@@ -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")
);
}
}
server/routers/aiBudget/updateAiBudget.ts
+52 -11
View File
@@ -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;
}
server/routers/aiBudget/validation.ts
+4 -2
View File
@@ -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"]
});
}
server/routers/aiGateway/pipeline.ts
+2 -1
View File
@@ -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;
server/routers/external.ts
+7
View File
@@ -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,