diff --git a/server/db/pg/schema/schema.ts b/server/db/pg/schema/schema.ts index 51552233c..6a24064a6 100644 --- a/server/db/pg/schema/schema.ts +++ b/server/db/pg/schema/schema.ts @@ -3,6 +3,7 @@ import { InferSelectModel, sql } from "drizzle-orm"; import { bigint, boolean, + check, index, integer, pgTable, @@ -214,7 +215,10 @@ export const resources = pgTable( aiProviderId: integer("aiProviderId").references( () => aiProviders.providerId, { onDelete: "set null" } - ) + ), + modelAccessMode: varchar("modelAccessMode").$type< + "passthrough" | "catalog" | "allowlist" + >() }, (t) => [ index("idx_resources_fulldomain") @@ -496,7 +500,10 @@ export const siteResources = pgTable( aiProviderId: integer("aiProviderId").references( () => aiProviders.providerId, { onDelete: "set null" } - ) + ), + modelAccessMode: varchar("modelAccessMode").$type< + "passthrough" | "catalog" | "allowlist" + >() }, (t) => [index("idx_siteresources_orgid_niceid").on(t.orgId, t.niceId)] ); @@ -1623,8 +1630,6 @@ export const aiProviders = pgTable("aiProviders", { skipTlsVerification: boolean("skipTlsVerification") .notNull() .default(false), - budgetAmount: real("budgetAmount"), - budgetUnit: varchar("budgetUnit").$type<"usd" | "tokens">(), enabled: boolean("enabled").notNull().default(true), createdAt: bigint("createdAt", { mode: "number" }).notNull(), updatedAt: bigint("updatedAt", { mode: "number" }).notNull() @@ -1639,8 +1644,6 @@ export const aiModels = pgTable( .references(() => aiProviders.providerId, { onDelete: "cascade" }), modelKey: varchar("modelKey").notNull(), name: varchar("name").notNull(), - budgetAmount: real("budgetAmount"), - budgetUnit: varchar("budgetUnit").$type<"usd" | "tokens">(), enabled: boolean("enabled").notNull().default(true), createdAt: bigint("createdAt", { mode: "number" }).notNull(), updatedAt: bigint("updatedAt", { mode: "number" }).notNull() @@ -1648,6 +1651,73 @@ export const aiModels = pgTable( (t) => [unique("ai_model_provider_key_uniq").on(t.providerId, t.modelKey)] ); +export const aiBudgets = pgTable( + "aiBudgets", + { + budgetId: serial("budgetId").primaryKey(), + orgId: varchar("orgId") + .notNull() + .references(() => orgs.orgId, { onDelete: "cascade" }), + providerId: integer("providerId").references( + () => aiProviders.providerId, + { onDelete: "cascade" } + ), + modelId: integer("modelId").references(() => aiModels.modelId, { + onDelete: "cascade" + }), + resourceId: integer("resourceId").references( + () => resources.resourceId, + { onDelete: "cascade" } + ), + siteResourceId: integer("siteResourceId").references( + () => siteResources.siteResourceId, + { onDelete: "cascade" } + ), + amount: real("amount").notNull(), + unit: varchar("unit").$type<"usd" | "tokens">().notNull(), + period: varchar("period") + .$type<"monthly">() + .notNull() + .default("monthly"), + enforcement: varchar("enforcement") + .$type<"hard" | "soft">() + .notNull() + .default("hard"), + enabled: boolean("enabled").notNull().default(true), + createdAt: bigint("createdAt", { mode: "number" }).notNull(), + updatedAt: bigint("updatedAt", { mode: "number" }).notNull() + }, + (t) => [ + check( + "ai_budget_one_scope", + sql`( + (CASE WHEN ${t.providerId} IS NOT NULL THEN 1 ELSE 0 END) + + (CASE WHEN ${t.modelId} IS NOT NULL THEN 1 ELSE 0 END) + + (CASE WHEN ${t.resourceId} IS NOT NULL THEN 1 ELSE 0 END) + + (CASE WHEN ${t.siteResourceId} IS NOT NULL THEN 1 ELSE 0 END) + ) = 1` + ), + 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) + ] +); + +export const aiBudgetPeriods = pgTable( + "aiBudgetPeriods", + { + periodId: serial("periodId").primaryKey(), + budgetId: integer("budgetId") + .notNull() + .references(() => aiBudgets.budgetId, { onDelete: "cascade" }), + periodStart: bigint("periodStart", { mode: "number" }).notNull(), + periodEnd: bigint("periodEnd", { mode: "number" }).notNull(), + usedAmount: real("usedAmount").notNull().default(0) + }, + (t) => [unique("ai_budget_period_start_uniq").on(t.budgetId, t.periodStart)] +); + export type Org = InferSelectModel; export type User = InferSelectModel; export type Site = InferSelectModel; @@ -1734,5 +1804,7 @@ export type UserPolicy = InferSelectModel; export type ResourcePolicyRule = InferSelectModel; export type AiProvider = InferSelectModel; export type AiModel = InferSelectModel; +export type AiBudget = InferSelectModel; +export type AiBudgetPeriod = InferSelectModel; export type ResourceAiModel = InferSelectModel; export type SiteResourceAiModel = InferSelectModel; diff --git a/server/db/sqlite/schema/schema.ts b/server/db/sqlite/schema/schema.ts index 3d3ff2ded..f453d36f7 100644 --- a/server/db/sqlite/schema/schema.ts +++ b/server/db/sqlite/schema/schema.ts @@ -1,6 +1,7 @@ import { randomUUID } from "crypto"; -import { InferSelectModel } from "drizzle-orm"; +import { InferSelectModel, sql } from "drizzle-orm"; import { + check, index, integer, primaryKey, @@ -219,7 +220,10 @@ export const resources = sqliteTable("resources", { aiProviderId: integer("aiProviderId").references( () => aiProviders.providerId, { onDelete: "set null" } - ) + ), + modelAccessMode: text("modelAccessMode").$type< + "passthrough" | "catalog" | "allowlist" + >() }); export const resourceAiModels = sqliteTable( @@ -483,7 +487,10 @@ export const siteResources = sqliteTable("siteResources", { aiProviderId: integer("aiProviderId").references( () => aiProviders.providerId, { onDelete: "set null" } - ) + ), + modelAccessMode: text("modelAccessMode").$type< + "passthrough" | "catalog" | "allowlist" + >() }); export const siteResourceAiModels = sqliteTable( @@ -1605,8 +1612,6 @@ export const aiProviders = sqliteTable("aiProviders", { skipTlsVerification: integer("skipTlsVerification", { mode: "boolean" }) .notNull() .default(false), - budgetAmount: real("budgetAmount"), - budgetUnit: text("budgetUnit").$type<"usd" | "tokens">(), enabled: integer("enabled", { mode: "boolean" }).notNull().default(true), createdAt: integer("createdAt").notNull(), updatedAt: integer("updatedAt").notNull() @@ -1621,8 +1626,6 @@ export const aiModels = sqliteTable( .references(() => aiProviders.providerId, { onDelete: "cascade" }), modelKey: text("modelKey").notNull(), name: text("name").notNull(), - budgetAmount: real("budgetAmount"), - budgetUnit: text("budgetUnit").$type<"usd" | "tokens">(), enabled: integer("enabled", { mode: "boolean" }) .notNull() .default(true), @@ -1632,6 +1635,72 @@ export const aiModels = sqliteTable( (t) => [unique("ai_model_provider_key_uniq").on(t.providerId, t.modelKey)] ); +export const aiBudgets = sqliteTable( + "aiBudgets", + { + budgetId: integer("budgetId").primaryKey({ autoIncrement: true }), + orgId: text("orgId") + .notNull() + .references(() => orgs.orgId, { onDelete: "cascade" }), + providerId: integer("providerId").references( + () => aiProviders.providerId, + { onDelete: "cascade" } + ), + modelId: integer("modelId").references(() => aiModels.modelId, { + onDelete: "cascade" + }), + resourceId: integer("resourceId").references( + () => resources.resourceId, + { onDelete: "cascade" } + ), + siteResourceId: integer("siteResourceId").references( + () => siteResources.siteResourceId, + { onDelete: "cascade" } + ), + amount: real("amount").notNull(), + unit: text("unit").$type<"usd" | "tokens">().notNull(), + period: text("period").$type<"monthly">().notNull().default("monthly"), + enforcement: text("enforcement") + .$type<"hard" | "soft">() + .notNull() + .default("hard"), + enabled: integer("enabled", { mode: "boolean" }) + .notNull() + .default(true), + createdAt: integer("createdAt").notNull(), + updatedAt: integer("updatedAt").notNull() + }, + (t) => [ + check( + "ai_budget_one_scope", + sql`( + (CASE WHEN ${t.providerId} IS NOT NULL THEN 1 ELSE 0 END) + + (CASE WHEN ${t.modelId} IS NOT NULL THEN 1 ELSE 0 END) + + (CASE WHEN ${t.resourceId} IS NOT NULL THEN 1 ELSE 0 END) + + (CASE WHEN ${t.siteResourceId} IS NOT NULL THEN 1 ELSE 0 END) + ) = 1` + ), + 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) + ] +); + +export const aiBudgetPeriods = sqliteTable( + "aiBudgetPeriods", + { + periodId: integer("periodId").primaryKey({ autoIncrement: true }), + budgetId: integer("budgetId") + .notNull() + .references(() => aiBudgets.budgetId, { onDelete: "cascade" }), + periodStart: integer("periodStart").notNull(), + periodEnd: integer("periodEnd").notNull(), + usedAmount: real("usedAmount").notNull().default(0) + }, + (t) => [unique("ai_budget_period_start_uniq").on(t.budgetId, t.periodStart)] +); + export type Org = InferSelectModel; export type User = InferSelectModel; export type Site = InferSelectModel; @@ -1716,5 +1785,7 @@ export type RolePolicy = InferSelectModel; export type UserPolicy = InferSelectModel; export type AiProvider = InferSelectModel; export type AiModel = InferSelectModel; +export type AiBudget = InferSelectModel; +export type AiBudgetPeriod = InferSelectModel; export type ResourceAiModel = InferSelectModel; export type SiteResourceAiModel = InferSelectModel;