diff --git a/server/db/pg/schema/schema.ts b/server/db/pg/schema/schema.ts index e2eb50205..7ded27eb7 100644 --- a/server/db/pg/schema/schema.ts +++ b/server/db/pg/schema/schema.ts @@ -1746,10 +1746,15 @@ export const aiBudgets = pgTable( updatedAt: bigint("updatedAt", { mode: "number" }).notNull() }, (t) => [ - unique("ai_budget_provider_uniq").on(t.providerId), - unique("ai_budget_model_uniq").on(t.modelId), - unique("ai_budget_resource_uniq").on(t.resourceId), - unique("ai_budget_site_resource_uniq").on(t.siteResourceId) + unique("ai_budget_provider_uniq").on(t.providerId, t.unit, t.period), + unique("ai_budget_model_uniq").on(t.modelId, t.unit, t.period), + unique("ai_budget_resource_uniq").on(t.resourceId, t.unit, t.period), + unique("ai_budget_site_resource_uniq").on( + t.siteResourceId, + t.unit, + t.period + ), + unique("ai_budget_role_uniq").on(t.roleId, t.unit, t.period) ] ); diff --git a/server/db/sqlite/schema/schema.ts b/server/db/sqlite/schema/schema.ts index c8d0042df..34fe67a0f 100644 --- a/server/db/sqlite/schema/schema.ts +++ b/server/db/sqlite/schema/schema.ts @@ -1732,10 +1732,15 @@ export const aiBudgets = sqliteTable( updatedAt: integer("updatedAt").notNull() }, (t) => [ - unique("ai_budget_provider_uniq").on(t.providerId), - unique("ai_budget_model_uniq").on(t.modelId), - unique("ai_budget_resource_uniq").on(t.resourceId), - unique("ai_budget_site_resource_uniq").on(t.siteResourceId) + unique("ai_budget_provider_uniq").on(t.providerId, t.unit, t.period), + unique("ai_budget_model_uniq").on(t.modelId, t.unit, t.period), + unique("ai_budget_resource_uniq").on(t.resourceId, t.unit, t.period), + unique("ai_budget_site_resource_uniq").on( + t.siteResourceId, + t.unit, + t.period + ), + unique("ai_budget_role_uniq").on(t.roleId, t.unit, t.period) ] ); diff --git a/server/routers/aiBudget/createAiBudget.ts b/server/routers/aiBudget/createAiBudget.ts index 705167c3a..3ddd76174 100644 --- a/server/routers/aiBudget/createAiBudget.ts +++ b/server/routers/aiBudget/createAiBudget.ts @@ -15,7 +15,7 @@ import createHttpError from "http-errors"; import logger from "@server/logger"; import { fromError } from "zod-validation-error"; import { OpenAPITags, registry } from "@server/openApi"; -import { eq } from "drizzle-orm"; +import { and, eq, isNull } from "drizzle-orm"; import type { CreateOrEditAiBudgetResponse } from "@server/routers/aiBudget/types"; import { aiBudgetEnforcementSchema, @@ -189,7 +189,7 @@ export async function createAiBudget( } } - const conflictCondition = + const scopeCondition = providerId !== undefined ? eq(aiBudgets.providerId, providerId) : modelId !== undefined @@ -198,22 +198,35 @@ export async function createAiBudget( ? eq(aiBudgets.resourceId, resourceId) : siteResourceId !== undefined ? eq(aiBudgets.siteResourceId, siteResourceId) - : undefined; + : 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) + ); - if (conflictCondition) { - const [existing] = await db - .select({ budgetId: aiBudgets.budgetId }) - .from(aiBudgets) - .where(conflictCondition) - .limit(1); - if (existing) { - return next( - createHttpError( - HttpCode.CONFLICT, - "A budget already exists for this scope" - ) - ); - } + const [existing] = await db + .select({ budgetId: aiBudgets.budgetId }) + .from(aiBudgets) + .where( + and( + scopeCondition, + eq(aiBudgets.unit, unit), + eq(aiBudgets.period, period) + ) + ) + .limit(1); + if (existing) { + return next( + createHttpError( + HttpCode.CONFLICT, + `A ${period} ${unit} budget already exists for this scope` + ) + ); } const now = Date.now(); diff --git a/server/routers/aiBudget/index.ts b/server/routers/aiBudget/index.ts index 0da53e0d2..36eed7ba4 100644 --- a/server/routers/aiBudget/index.ts +++ b/server/routers/aiBudget/index.ts @@ -1,5 +1,10 @@ export * from "./createAiBudget"; export * from "./listAiBudgets"; +export * from "./listAiBudgetsForProvider"; +export * from "./listAiBudgetsForModel"; +export * from "./listAiBudgetsForResource"; +export * from "./listAiBudgetsForSiteResource"; +export * from "./listAiBudgetsForRole"; export * from "./getAiBudget"; export * from "./updateAiBudget"; export * from "./deleteAiBudget"; diff --git a/server/routers/aiBudget/listAiBudgetsForModel.ts b/server/routers/aiBudget/listAiBudgetsForModel.ts new file mode 100644 index 000000000..4a2277b5b --- /dev/null +++ b/server/routers/aiBudget/listAiBudgetsForModel.ts @@ -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({ + modelId: z.coerce.number().int().positive() +}); + +registry.registerPath({ + method: "get", + path: "/ai-model/{modelId}/ai-budgets", + description: "List AI budgets scoped to an AI model.", + tags: [OpenAPITags.AiBudget], + request: { + params: paramsSchema + }, + responses: { + 200: { + description: "Successful response" + } + } +}); + +export async function listAiBudgetsForModel( + req: Request, + res: Response, + next: NextFunction +): Promise { + try { + const parsedParams = paramsSchema.safeParse(req.params); + if (!parsedParams.success) { + return next( + createHttpError( + HttpCode.BAD_REQUEST, + fromError(parsedParams.error).toString() + ) + ); + } + + const { modelId } = parsedParams.data; + + const budgets = await db + .select() + .from(aiBudgets) + .where(eq(aiBudgets.modelId, modelId)) + .orderBy(asc(aiBudgets.unit), asc(aiBudgets.period)); + + return response(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") + ); + } +} diff --git a/server/routers/aiBudget/listAiBudgetsForProvider.ts b/server/routers/aiBudget/listAiBudgetsForProvider.ts new file mode 100644 index 000000000..1a0625f4b --- /dev/null +++ b/server/routers/aiBudget/listAiBudgetsForProvider.ts @@ -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({ + providerId: z.coerce.number().int().positive() +}); + +registry.registerPath({ + method: "get", + path: "/ai-provider/{providerId}/ai-budgets", + description: "List AI budgets scoped to an AI provider.", + tags: [OpenAPITags.AiBudget], + request: { + params: paramsSchema + }, + responses: { + 200: { + description: "Successful response" + } + } +}); + +export async function listAiBudgetsForProvider( + req: Request, + res: Response, + next: NextFunction +): Promise { + try { + const parsedParams = paramsSchema.safeParse(req.params); + if (!parsedParams.success) { + return next( + createHttpError( + HttpCode.BAD_REQUEST, + fromError(parsedParams.error).toString() + ) + ); + } + + const { providerId } = parsedParams.data; + + const budgets = await db + .select() + .from(aiBudgets) + .where(eq(aiBudgets.providerId, providerId)) + .orderBy(asc(aiBudgets.unit), asc(aiBudgets.period)); + + return response(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") + ); + } +} diff --git a/server/routers/aiBudget/listAiBudgetsForResource.ts b/server/routers/aiBudget/listAiBudgetsForResource.ts new file mode 100644 index 000000000..106b049a7 --- /dev/null +++ b/server/routers/aiBudget/listAiBudgetsForResource.ts @@ -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({ + resourceId: z.coerce.number().int().positive() +}); + +registry.registerPath({ + method: "get", + path: "/resource/{resourceId}/ai-budgets", + description: "List AI budgets scoped to a resource.", + tags: [OpenAPITags.AiBudget], + request: { + params: paramsSchema + }, + responses: { + 200: { + description: "Successful response" + } + } +}); + +export async function listAiBudgetsForResource( + req: Request, + res: Response, + next: NextFunction +): Promise { + try { + const parsedParams = paramsSchema.safeParse(req.params); + if (!parsedParams.success) { + return next( + createHttpError( + HttpCode.BAD_REQUEST, + fromError(parsedParams.error).toString() + ) + ); + } + + const { resourceId } = parsedParams.data; + + const budgets = await db + .select() + .from(aiBudgets) + .where(eq(aiBudgets.resourceId, resourceId)) + .orderBy(asc(aiBudgets.unit), asc(aiBudgets.period)); + + return response(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") + ); + } +} diff --git a/server/routers/aiBudget/listAiBudgetsForRole.ts b/server/routers/aiBudget/listAiBudgetsForRole.ts new file mode 100644 index 000000000..2e8ff0e8d --- /dev/null +++ b/server/routers/aiBudget/listAiBudgetsForRole.ts @@ -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({ + roleId: z.coerce.number().int().positive() +}); + +registry.registerPath({ + method: "get", + path: "/role/{roleId}/ai-budgets", + description: "List AI budgets scoped to a role.", + tags: [OpenAPITags.AiBudget], + request: { + params: paramsSchema + }, + responses: { + 200: { + description: "Successful response" + } + } +}); + +export async function listAiBudgetsForRole( + req: Request, + res: Response, + next: NextFunction +): Promise { + try { + const parsedParams = paramsSchema.safeParse(req.params); + if (!parsedParams.success) { + return next( + createHttpError( + HttpCode.BAD_REQUEST, + fromError(parsedParams.error).toString() + ) + ); + } + + const { roleId } = parsedParams.data; + + const budgets = await db + .select() + .from(aiBudgets) + .where(eq(aiBudgets.roleId, roleId)) + .orderBy(asc(aiBudgets.unit), asc(aiBudgets.period)); + + return response(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") + ); + } +} diff --git a/server/routers/aiBudget/listAiBudgetsForSiteResource.ts b/server/routers/aiBudget/listAiBudgetsForSiteResource.ts new file mode 100644 index 000000000..b2557719d --- /dev/null +++ b/server/routers/aiBudget/listAiBudgetsForSiteResource.ts @@ -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({ + siteResourceId: z.coerce.number().int().positive() +}); + +registry.registerPath({ + method: "get", + path: "/site-resource/{siteResourceId}/ai-budgets", + description: "List AI budgets scoped to a site resource.", + tags: [OpenAPITags.AiBudget], + request: { + params: paramsSchema + }, + responses: { + 200: { + description: "Successful response" + } + } +}); + +export async function listAiBudgetsForSiteResource( + req: Request, + res: Response, + next: NextFunction +): Promise { + try { + const parsedParams = paramsSchema.safeParse(req.params); + if (!parsedParams.success) { + return next( + createHttpError( + HttpCode.BAD_REQUEST, + fromError(parsedParams.error).toString() + ) + ); + } + + const { siteResourceId } = parsedParams.data; + + const budgets = await db + .select() + .from(aiBudgets) + .where(eq(aiBudgets.siteResourceId, siteResourceId)) + .orderBy(asc(aiBudgets.unit), asc(aiBudgets.period)); + + return response(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") + ); + } +} diff --git a/server/routers/aiBudget/types.ts b/server/routers/aiBudget/types.ts index 11ffd1050..c689ef30e 100644 --- a/server/routers/aiBudget/types.ts +++ b/server/routers/aiBudget/types.ts @@ -5,6 +5,10 @@ export type ListAiBudgetsResponse = PaginatedResponse<{ budgets: AiBudget[]; }>; +export type ListAiBudgetsByScopeResponse = { + budgets: AiBudget[]; +}; + export type GetAiBudgetResponse = { budget: AiBudget; }; diff --git a/server/routers/aiBudget/updateAiBudget.ts b/server/routers/aiBudget/updateAiBudget.ts index e8bd845c5..723b346c7 100644 --- a/server/routers/aiBudget/updateAiBudget.ts +++ b/server/routers/aiBudget/updateAiBudget.ts @@ -15,7 +15,7 @@ import createHttpError from "http-errors"; import logger from "@server/logger"; import { fromError } from "zod-validation-error"; import { OpenAPITags, registry } from "@server/openApi"; -import { and, eq, ne } from "drizzle-orm"; +import { and, eq, isNull, ne } from "drizzle-orm"; import type { CreateOrEditAiBudgetResponse } from "@server/routers/aiBudget/types"; import { aiBudgetEnforcementSchema, @@ -128,6 +128,9 @@ export async function updateAiBudget( : existing.siteResourceId; const nextRoleId = body.roleId !== undefined ? body.roleId : existing.roleId; + const nextUnit = body.unit !== undefined ? body.unit : existing.unit; + const nextPeriod = + body.period !== undefined ? body.period : existing.period; const scopeValidation = z .object({ @@ -242,44 +245,45 @@ export async function updateAiBudget( } } - const conflictCondition = - body.providerId !== undefined && body.providerId !== null - ? and( - eq(aiBudgets.providerId, body.providerId), - ne(aiBudgets.budgetId, budgetId) - ) - : body.modelId !== undefined && body.modelId !== null - ? and( - eq(aiBudgets.modelId, body.modelId), - ne(aiBudgets.budgetId, budgetId) - ) - : body.resourceId !== undefined && body.resourceId !== null - ? and( - eq(aiBudgets.resourceId, body.resourceId), - ne(aiBudgets.budgetId, budgetId) - ) - : body.siteResourceId !== undefined && - body.siteResourceId !== null - ? and( - eq(aiBudgets.siteResourceId, body.siteResourceId), - ne(aiBudgets.budgetId, budgetId) - ) - : undefined; + const scopeCondition = + nextProviderId !== null + ? eq(aiBudgets.providerId, nextProviderId) + : nextModelId !== null + ? eq(aiBudgets.modelId, nextModelId) + : nextResourceId !== null + ? eq(aiBudgets.resourceId, nextResourceId) + : nextSiteResourceId !== null + ? 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) + ); - if (conflictCondition) { - const [conflict] = await db - .select({ budgetId: aiBudgets.budgetId }) - .from(aiBudgets) - .where(conflictCondition) - .limit(1); - if (conflict) { - return next( - createHttpError( - HttpCode.CONFLICT, - "A budget already exists for this scope" - ) - ); - } + const [conflict] = await db + .select({ budgetId: aiBudgets.budgetId }) + .from(aiBudgets) + .where( + and( + scopeCondition, + eq(aiBudgets.unit, nextUnit), + eq(aiBudgets.period, nextPeriod), + ne(aiBudgets.budgetId, budgetId) + ) + ) + .limit(1); + if (conflict) { + return next( + createHttpError( + HttpCode.CONFLICT, + `A ${nextPeriod} ${nextUnit} budget already exists for this scope` + ) + ); } const updateData: Partial = { diff --git a/server/routers/external.ts b/server/routers/external.ts index f469707f8..41df6ea59 100644 --- a/server/routers/external.ts +++ b/server/routers/external.ts @@ -1626,6 +1626,41 @@ authenticated.delete( aiBudget.deleteAiBudget ); +authenticated.get( + "/ai-provider/:providerId/ai-budgets", + verifyAiProviderAccess, + verifyUserHasAction(ActionsEnum.listAiBudgets), + aiBudget.listAiBudgetsForProvider +); + +authenticated.get( + "/ai-model/:modelId/ai-budgets", + verifyAiModelAccess, + verifyUserHasAction(ActionsEnum.listAiBudgets), + aiBudget.listAiBudgetsForModel +); + +authenticated.get( + "/resource/:resourceId/ai-budgets", + verifyResourceAccess, + verifyUserHasAction(ActionsEnum.listAiBudgets), + aiBudget.listAiBudgetsForResource +); + +authenticated.get( + "/site-resource/:siteResourceId/ai-budgets", + verifySiteResourceAccess, + verifyUserHasAction(ActionsEnum.listAiBudgets), + aiBudget.listAiBudgetsForSiteResource +); + +authenticated.get( + "/role/:roleId/ai-budgets", + verifyRoleAccess, + verifyUserHasAction(ActionsEnum.listAiBudgets), + aiBudget.listAiBudgetsForRole +); + authenticated.get( "/org/:orgId/labels", verifyOrgAccess,