plugin-ai-api 1.0.25 → 1.1.0

package/src/server/billing.ts

@@ -5,6 +5,8 @@ import timezone from 'dayjs/plugin/timezone';
 import type { Model } from '@nocobase/database';
 import type { Transaction } from 'sequelize';
 import type { Usage } from './usage';
+import { type QuotaMode } from './quota-groups';
+import { getAiApiConfig, resolveRequestUserGroup } from './utils/request-cache';
 
 dayjs.extend(utc);
 dayjs.extend(timezone);
@@ -28,7 +30,8 @@ export interface PriceSnapshot {
 
 interface QuotaReservation {
   bucketId: string | number | bigint;
-  policyId: string | number | bigint;
+  groupId: string | number | bigint;
+  quotaMode: QuotaMode;
   estimatedInputTokens: number;
   estimatedOutputTokens: number;
   reservedTokens: number;
@@ -54,6 +57,8 @@ export interface BillingFinalization {
   costStatus?: 'calculated' | 'estimated' | 'unpriced' | 'usage_unavailable';
   modelPriceId?: string | number | bigint;
   quotaPolicyId?: string | number | bigint;
+  groupId?: string | number | bigint;
+  quotaMode?: QuotaMode;
   inputPricePerMillionTokens?: string;
   outputPricePerMillionTokens?: string;
   fixedCostPerRequest?: string;
@@ -195,27 +200,23 @@ export async function prepareLlmBilling(ctx: Context, resolved: ResolvedLlmModel
   };
   stateOf(ctx).aiApiLlmBilling = billing;
 
-  const config = await ctx.db.getRepository('aiApiConfig').findOne();
+  const config = await getAiApiConfig(ctx);
   if (!valueOf<boolean | undefined>(config, 'quotaEnabled') || userId === undefined || userId === null) return;
 
-  const policy = await ctx.db.getRepository('aiApiUserQuotaPolicies').findOne({
-    filter: { userId, enabled: true },
-    sort: '-updatedAt',
-  });
-  if (!policy) return;
+  const group = await resolveRequestUserGroup(ctx, userId);
+  if (!group.enabled) return;
 
-  const rejectUnpriced = valueOf<boolean>(policy, 'rejectUnpricedModel');
+  const rejectUnpriced = group.rejectUnpricedModel;
   if (!price && rejectUnpriced) {
     throw new AiApiQuotaError(
       'model_price_not_configured',
       `Pricing is not configured for '${serviceName}/${resolved.modelId}'.`,
     );
   }
-  const policyCurrency = valueOf<string>(policy, 'currency');
-  if (price && policyCurrency !== price.currency) {
+  if (price && group.currency !== price.currency) {
     throw new AiApiQuotaError(
       'quota_currency_mismatch',
-      `Quota currency '${policyCurrency}' does not match model price currency '${price.currency}'.`,
+      `Quota currency '${group.currency}' does not match model price currency '${price.currency}'.`,
     );
   }
 
@@ -225,14 +226,16 @@ export async function prepareLlmBilling(ctx: Context, resolved: ResolvedLlmModel
   const estimatedOutputTokens = normalizePositiveInteger(body.max_completion_tokens ?? body.max_tokens, defaultOutput);
   const reservedTokens = estimatedInputTokens + estimatedOutputTokens;
   const reservedCost = price ? calculateLlmCost(estimatedInputTokens, estimatedOutputTokens, price) : '0.00000000';
-  const period = getPeriodBounds(valueOf<string>(policy, 'periodType'), valueOf<string>(policy, 'timezone'));
-  const Bucket = ctx.db.getModel('aiApiUserQuotaBuckets');
+  const period = getPeriodBounds(group.periodType, group.timezone);
+  const bucketUserId = group.quotaMode === 'share' ? 0 : userId;
+  const Bucket = ctx.db.getModel('aiApiGroupQuotaBuckets');
 
   const reservation = await ctx.db.sequelize.transaction(async (transaction: Transaction) => {
     const [bucket] = await Bucket.findOrCreate({
-      where: { policyId: valueOf(policy, 'id'), periodStart: period.start },
+      where: { groupId: group.id, userId: bucketUserId, periodStart: period.start },
       defaults: {
-        userId,
+        groupId: group.id,
+        userId: bucketUserId,
         periodEnd: period.end,
         requestCount: 0,
         totalTokens: 0,
@@ -247,32 +250,20 @@ export async function prepareLlmBilling(ctx: Context, resolved: ResolvedLlmModel
 
     const requestCount = BigInt(String(bucket.get('requestCount') ?? 0));
     const reservedRequests = BigInt(String(bucket.get('reservedRequests') ?? 0));
-    if (exceedsIntegerLimit(requestCount + reservedRequests, 1n, valueOf(policy, 'requestLimit'))) {
-      throw new AiApiQuotaError('request_quota_exceeded', 'The request quota for this user has been exceeded.');
+    if (exceedsIntegerLimit(requestCount + reservedRequests, 1n, group.requestLimit)) {
+      throw new AiApiQuotaError('request_quota_exceeded', 'The request quota for this group has been exceeded.');
     }
 
     const totalTokens = BigInt(String(bucket.get('totalTokens') ?? 0));
     const alreadyReservedTokens = BigInt(String(bucket.get('reservedTokens') ?? 0));
-    if (
-      exceedsIntegerLimit(
-        totalTokens + alreadyReservedTokens,
-        BigInt(reservedTokens),
-        valueOf(policy, 'totalTokenLimit'),
-      )
-    ) {
-      throw new AiApiQuotaError('token_quota_exceeded', 'The token quota for this user has been exceeded.');
+    if (exceedsIntegerLimit(totalTokens + alreadyReservedTokens, BigInt(reservedTokens), group.totalTokenLimit)) {
+      throw new AiApiQuotaError('token_quota_exceeded', 'The token quota for this group has been exceeded.');
     }
 
     const cost = decimalUnits(bucket.get('cost'), COST_SCALE);
     const alreadyReservedCost = decimalUnits(bucket.get('reservedCost'), COST_SCALE);
-    if (
-      exceedsDecimalLimit(
-        cost + alreadyReservedCost,
-        decimalUnits(reservedCost, COST_SCALE),
-        valueOf(policy, 'costLimit'),
-      )
-    ) {
-      throw new AiApiQuotaError('cost_quota_exceeded', 'The cost quota for this user has been exceeded.');
+    if (exceedsDecimalLimit(cost + alreadyReservedCost, decimalUnits(reservedCost, COST_SCALE), group.costLimit)) {
+      throw new AiApiQuotaError('cost_quota_exceeded', 'The cost quota for this group has been exceeded.');
     }
 
     await bucket.update(
@@ -285,13 +276,13 @@ export async function prepareLlmBilling(ctx: Context, resolved: ResolvedLlmModel
     );
     return {
       bucketId: bucket.get('id') as string | number | bigint,
-      policyId: valueOf<string | number | bigint>(policy, 'id'),
+      groupId: group.id,
+      quotaMode: group.quotaMode,
       estimatedInputTokens,
       estimatedOutputTokens,
       reservedTokens,
       reservedCost,
-      missingUsageBehavior:
-        valueOf<string>(policy, 'missingUsageBehavior') === 'allow' ? ('allow' as const) : ('use_reserved' as const),
+      missingUsageBehavior: group.missingUsageBehavior === 'allow' ? ('allow' as const) : ('use_reserved' as const),
     };
   });
   billing.reservation = reservation;
@@ -337,7 +328,7 @@ export async function finalizeLlmBilling(
   const cost = numbers && billing.price ? calculateLlmCost(numbers.input, numbers.output, billing.price) : undefined;
   const reservation = billing.reservation;
   if (reservation) {
-    const Bucket = ctx.db.getModel('aiApiUserQuotaBuckets');
+    const Bucket = ctx.db.getModel('aiApiGroupQuotaBuckets');
     await ctx.db.sequelize.transaction(async (transaction: Transaction) => {
       const bucket = await Bucket.findByPk(reservation.bucketId, { transaction, lock: transaction.LOCK.UPDATE });
       if (!bucket) return;
@@ -384,7 +375,8 @@ export async function finalizeLlmBilling(
     currency: billing.price?.currency,
     costStatus,
     modelPriceId: billing.price?.id,
-    quotaPolicyId: reservation?.policyId,
+    groupId: reservation?.groupId,
+    quotaMode: reservation?.quotaMode,
     inputPricePerMillionTokens: billing.price?.inputPricePerMillionTokens,
     outputPricePerMillionTokens: billing.price?.outputPricePerMillionTokens,
     fixedCostPerRequest: billing.price?.fixedCostPerRequest,

package/src/server/collections/ai-api-config.ts

@@ -22,7 +22,7 @@ export default defineCollection({
     {
       name: 'defaultAiEmployee',
       type: 'string',
-      comment: 'Username of the default AI Employee for system prompt injection',
+      comment: 'Username of the default AI Employee used by agent mode. Direct LLM mode ignores it.',
     },
     {
       name: 'defaultLlmService',
@@ -35,12 +35,6 @@ export default defineCollection({
       defaultValue: [],
       comment: 'Array of llmService names to expose. Empty = expose all enabled services',
     },
-    {
-      name: 'rateLimitPerMinute',
-      type: 'integer',
-      defaultValue: 60,
-      comment: 'Max requests per user per minute',
-    },
     {
       name: 'maxRequestBodyMb',
       type: 'integer',

package/src/server/collections/ai-api-group-members.ts

@@ -0,0 +1,41 @@
+/**
+ * This file is part of the NocoBase (R) project.
+ * Copyright (c) 2020-2024 NocoBase Co., Ltd.
+ * Authors: NocoBase Team.
+ *
+ * This project is dual-licensed under AGPL-3.0 and NocoBase Commercial License.
+ * For more information, please refer to: https://www.nocobase.com/agreement.
+ */
+
+import { defineCollection } from '@nocobase/database';
+
+export default defineCollection({
+  name: 'aiApiGroupMembers',
+  autoGenId: true,
+  fields: [
+    { name: 'groupId', type: 'bigInt', allowNull: false, index: true },
+    {
+      name: 'group',
+      type: 'belongsTo',
+      target: 'aiApiUsageGroups',
+      targetKey: 'id',
+      foreignKey: 'groupId',
+      constraints: false,
+    },
+    { name: 'userId', type: 'bigInt', allowNull: false, index: true },
+    {
+      name: 'user',
+      type: 'belongsTo',
+      target: 'users',
+      targetKey: 'id',
+      foreignKey: 'userId',
+      constraints: false,
+    },
+  ],
+  indexes: [
+    {
+      fields: ['userId'],
+      unique: true,
+    },
+  ],
+});

package/src/server/collections/ai-api-group-quota-buckets.ts

@@ -0,0 +1,42 @@
+/**
+ * This file is part of the NocoBase (R) project.
+ * Copyright (c) 2020-2024 NocoBase Co., Ltd.
+ * Authors: NocoBase Team.
+ *
+ * This project is dual-licensed under AGPL-3.0 and NocoBase Commercial License.
+ * For more information, please refer to: https://www.nocobase.com/agreement.
+ */
+
+import { defineCollection } from '@nocobase/database';
+
+export default defineCollection({
+  name: 'aiApiGroupQuotaBuckets',
+  autoGenId: true,
+  fields: [
+    { name: 'groupId', type: 'bigInt', allowNull: false, index: true },
+    {
+      name: 'group',
+      type: 'belongsTo',
+      target: 'aiApiUsageGroups',
+      targetKey: 'id',
+      foreignKey: 'groupId',
+      constraints: false,
+    },
+    // userId = 0 means the shared bucket in share mode; real user ids are always > 0.
+    { name: 'userId', type: 'bigInt', allowNull: false, defaultValue: 0, index: true },
+    { name: 'periodStart', type: 'datetimeTz', allowNull: false, index: true },
+    { name: 'periodEnd', type: 'datetimeTz', allowNull: false },
+    { name: 'requestCount', type: 'bigInt', allowNull: false, defaultValue: 0 },
+    { name: 'totalTokens', type: 'bigInt', allowNull: false, defaultValue: 0 },
+    { name: 'cost', type: 'decimal', precision: 20, scale: 8, allowNull: false, defaultValue: 0 },
+    { name: 'reservedRequests', type: 'bigInt', allowNull: false, defaultValue: 0 },
+    { name: 'reservedTokens', type: 'bigInt', allowNull: false, defaultValue: 0 },
+    { name: 'reservedCost', type: 'decimal', precision: 20, scale: 8, allowNull: false, defaultValue: 0 },
+  ],
+  indexes: [
+    {
+      fields: ['groupId', 'userId', 'periodStart'],
+      unique: true,
+    },
+  ],
+});

package/src/server/collections/ai-api-model-metadata.ts

@@ -55,6 +55,13 @@ export default defineCollection({
       allowNull: true,
       comment: 'Human-readable description returned as description in the model object.',
     },
+    {
+      name: 'systemPrompt',
+      type: 'text',
+      allowNull: true,
+      comment:
+        'Initial system prompt prepended as the first system message of every request for this model. Never replaces the client system prompt.',
+    },
     {
       name: 'enabled',
       type: 'boolean',

package/src/server/collections/ai-api-usage-groups.ts

@@ -0,0 +1,53 @@
+/**
+ * This file is part of the NocoBase (R) project.
+ * Copyright (c) 2020-2024 NocoBase Co., Ltd.
+ * Authors: NocoBase Team.
+ *
+ * This project is dual-licensed under AGPL-3.0 and NocoBase Commercial License.
+ * For more information, please refer to: https://www.nocobase.com/agreement.
+ */
+
+import { defineCollection } from '@nocobase/database';
+
+export default defineCollection({
+  name: 'aiApiUsageGroups',
+  autoGenId: true,
+  fields: [
+    { name: 'name', type: 'string', allowNull: false },
+    { name: 'isDefault', type: 'boolean', defaultValue: false, index: true },
+    { name: 'quotaMode', type: 'string', allowNull: false, defaultValue: 'per_user' },
+    { name: 'rateLimitPerMinute', type: 'integer', allowNull: false, defaultValue: 60 },
+    { name: 'enabled', type: 'boolean', defaultValue: true, index: true },
+    { name: 'periodType', type: 'string', allowNull: false, defaultValue: 'monthly' },
+    { name: 'timezone', type: 'string', allowNull: false, defaultValue: 'UTC' },
+    { name: 'requestLimit', type: 'bigInt', allowNull: true },
+    { name: 'totalTokenLimit', type: 'bigInt', allowNull: true },
+    { name: 'costLimit', type: 'decimal', precision: 20, scale: 8, allowNull: true },
+    { name: 'currency', type: 'string', allowNull: false, defaultValue: 'USD' },
+    { name: 'rejectUnpricedModel', type: 'boolean', defaultValue: true },
+    { name: 'missingUsageBehavior', type: 'string', allowNull: false, defaultValue: 'use_reserved' },
+    { name: 'contextOverflowBehavior', type: 'string', allowNull: false, defaultValue: 'reject' },
+    // Model access: empty lists mean "no narrowing" — the group inherits the full
+    // global configuration. Non-empty lists narrow what members may use.
+    {
+      name: 'allowedLlmServices',
+      type: 'json',
+      defaultValue: [],
+      comment: 'Empty means all globally enabled services; otherwise narrows to these services.',
+    },
+    { name: 'allowAllModels', type: 'boolean', defaultValue: true },
+    {
+      name: 'allowedModels',
+      type: 'json',
+      defaultValue: [],
+      comment: 'Array of "serviceName/modelId" members may use (when allowAllModels=false).',
+    },
+  ],
+  indexes: [
+    {
+      fields: ['isDefault'],
+      unique: true,
+      where: { isDefault: true },
+    },
+  ],
+});

package/src/server/collections/ai-api-usage-records.ts

@@ -38,6 +38,7 @@ export default defineCollection({
     { name: 'costStatus', type: 'string', allowNull: true, index: true },
     { name: 'modelPriceId', type: 'bigInt', allowNull: true, index: true },
     { name: 'quotaPolicyId', type: 'bigInt', allowNull: true, index: true },
+    { name: 'groupId', type: 'bigInt', allowNull: true, index: true },
     { name: 'inputPricePerMillionTokens', type: 'decimal', allowNull: true, precision: 20, scale: 10 },
     { name: 'outputPricePerMillionTokens', type: 'decimal', allowNull: true, precision: 20, scale: 10 },
     { name: 'fixedCostPerRequest', type: 'decimal', allowNull: true, precision: 20, scale: 10 },

package/src/server/middleware/rate-limit.ts

@@ -10,14 +10,14 @@
 import { Context } from '@nocobase/actions';
 import { RateLimiter } from '../utils/rate-limiter';
 import { toOpenAIError } from '../utils/openai-format';
+import { resolveRequestUserGroup } from '../utils/request-cache';
 
 /**
  * Creates a rate-limiting check function for use in the AI API router.
  *
  * Must be called AFTER authenticateBearer() so ctx.state.currentUser is set.
- * Reads rateLimitPerMinute from aiApiConfig on each request (cheap single-row read,
- * allows config changes to take effect immediately without restart).
- * Falls back to 60 req/min if the config record is missing or the field is 0/null.
+ * Resolves the user's group and reads rateLimitPerMinute from that group.
+ * Falls back to 60 req/min if the default group is missing or misconfigured.
  *
  * Returns false (and writes the 429 response) when the rate limit is exceeded.
  * Returns true when the request is allowed.
@@ -35,21 +35,19 @@ export function createRateLimitMiddleware(limiter: RateLimiter) {
 
     let limit = 60;
     try {
-      const config = await ctx.db.getRepository('aiApiConfig').findOne();
-      const configLimit = config?.rateLimitPerMinute;
-      if (configLimit && configLimit > 0) {
-        limit = configLimit;
+      const group = await resolveRequestUserGroup(ctx, userId);
+      const groupLimit = group?.rateLimitPerMinute;
+      if (groupLimit && groupLimit > 0) {
+        limit = groupLimit;
       }
-    } catch (configErr) {
-      // Config read failure: fail open — don't block legitimate requests
-      // Log at WARN so admins can detect DB connectivity issues
-      ctx.app?.logger?.warn('[ai-api] Rate limit config read failed, using default (60/min)', configErr);
+    } catch (err) {
+      ctx.app?.logger?.warn('[ai-api] Rate limit group resolution failed, using default (60/min)', err);
     }
 
     const result = limiter.check(userId, limit);
 
     if (!result.allowed) {
-      const retryAfterSec = Math.ceil((result as any).retryAfterMs / 1000);
+      const retryAfterSec = Math.ceil(result.retryAfterMs / 1000);
       ctx.set('Retry-After', String(retryAfterSec));
       ctx.set('X-RateLimit-Limit', String(limit));
       ctx.set('X-RateLimit-Remaining', '0');

package/src/server/migrations/20260815000000-add-usage-groups.ts

@@ -0,0 +1,147 @@
+/**
+ * This file is part of the NocoBase (R) project.
+ * Copyright (c) 2020-2024 NocoBase Co., Ltd.
+ * Authors: NocoBase Team.
+ *
+ * This project is dual-licensed under AGPL-3.0 and NocoBase Commercial License.
+ * For more information, please refer to: https://www.nocobase.com/agreement.
+ */
+
+import { Migration } from '@nocobase/server';
+
+export default class AddUsageGroups extends Migration {
+  // Must run afterSync: plugin collections are only registered once the app has
+  // loaded, and the new group tables must exist before seeding.
+  on = 'afterSync' as const;
+
+  async up() {
+    const groupCollection = this.getGroupCollection();
+
+    if (!groupCollection) {
+      throw new Error('AI API usage groups migration could not resolve the aiApiUsageGroups collection.');
+    }
+
+    const configCollection = this.db.getCollection('aiApiConfig');
+    let rateLimitPerMinute = 60;
+    let quotaEnabled = false;
+    if (configCollection && (await configCollection.existsInDb())) {
+      // rateLimitPerMinute is no longer part of the collection definition, so
+      // read the legacy column directly before dropping it.
+      const tableName = configCollection.getTableNameWithSchema();
+      const columns = await this.queryInterface.describeTable(tableName);
+      if (columns.rateLimitPerMinute) {
+        const legacyRateLimit = await this.readLegacyRateLimit(configCollection.quotedTableName());
+        if (legacyRateLimit && legacyRateLimit > 0) {
+          rateLimitPerMinute = legacyRateLimit;
+        }
+      }
+      const config = await this.db.getRepository('aiApiConfig').findOne();
+      if (config) {
+        quotaEnabled = !!config.get('quotaEnabled');
+      }
+    }
+
+    const existingDefault = await this.db.getRepository('aiApiUsageGroups').findOne({
+      filter: { isDefault: true },
+    });
+
+    if (!existingDefault) {
+      await this.db.getRepository('aiApiUsageGroups').create({
+        values: {
+          name: 'Default',
+          isDefault: true,
+          quotaMode: 'per_user',
+          rateLimitPerMinute,
+          enabled: quotaEnabled,
+          periodType: 'monthly',
+          timezone: 'UTC',
+          requestLimit: null,
+          totalTokenLimit: null,
+          costLimit: null,
+          currency: 'USD',
+          rejectUnpricedModel: true,
+          missingUsageBehavior: 'use_reserved',
+          contextOverflowBehavior: 'reject',
+        },
+      });
+    }
+
+    if (configCollection && (await configCollection.existsInDb())) {
+      const tableName = configCollection.getTableNameWithSchema();
+      const columns = await this.queryInterface.describeTable(tableName);
+      if (columns.rateLimitPerMinute) {
+        await this.queryInterface.removeColumn(tableName, 'rateLimitPerMinute');
+      }
+    }
+
+    await this.migrateLegacyQuotaPolicies();
+  }
+
+  async down() {
+    // Recreate the rateLimitPerMinute column on aiApiConfig with a sensible default.
+    const configCollection = this.db.getCollection('aiApiConfig');
+    if (configCollection && (await configCollection.existsInDb())) {
+      const tableName = configCollection.getTableNameWithSchema();
+      const columns = await this.queryInterface.describeTable(tableName);
+      if (!columns.rateLimitPerMinute) {
+        await this.queryInterface.addColumn(tableName, 'rateLimitPerMinute', {
+          type: 'INTEGER',
+          allowNull: false,
+          defaultValue: 60,
+        });
+      }
+    }
+  }
+
+  private getGroupCollection() {
+    return this.db.getCollection('aiApiUsageGroups');
+  }
+
+  private async readLegacyRateLimit(quotedTable: string): Promise<number | undefined> {
+    const qi = this.db.sequelize.getQueryInterface();
+    const [rows] = await this.db.sequelize.query(
+      `SELECT ${qi.quoteIdentifier('rateLimitPerMinute')} FROM ${quotedTable} LIMIT 1`,
+    );
+    const first = (rows as Array<Record<string, unknown>> | undefined)?.[0];
+    const value = Number(first?.rateLimitPerMinute);
+    return Number.isFinite(value) ? value : undefined;
+  }
+
+  private async migrateLegacyQuotaPolicies() {
+    const defaultGroup = await this.db.getRepository('aiApiUsageGroups').findOne({
+      filter: { isDefault: true },
+    });
+    if (!defaultGroup) return;
+
+    const policyCollection = this.db.getCollection('aiApiUserQuotaPolicies');
+    const memberCollection = this.db.getCollection('aiApiGroupMembers');
+    const bucketCollection = this.db.getCollection('aiApiGroupQuotaBuckets');
+    if (!policyCollection || !memberCollection || !bucketCollection) return;
+
+    const policies = (await this.db.getRepository('aiApiUserQuotaPolicies').find({
+      pageSize: 1000,
+    })) as unknown[];
+
+    for (const policy of policies) {
+      const userId = this.valueOf(policy, 'userId');
+      if (userId === undefined || userId === null) continue;
+
+      const existingMember = await this.db.getRepository('aiApiGroupMembers').findOne({
+        filter: { userId },
+      });
+      if (!existingMember) {
+        await this.db.getRepository('aiApiGroupMembers').create({
+          values: { groupId: this.valueOf(defaultGroup, 'id'), userId },
+        });
+      }
+    }
+  }
+
+  private valueOf(model: unknown, name: string): unknown {
+    if (!model) return undefined;
+    if (typeof (model as Record<string, unknown>).get === 'function') {
+      return (model as Record<string, unknown>).get(name);
+    }
+    return (model as Record<string, unknown>)[name];
+  }
+}