Initial commit: OmniAI backend server

src/enterpriseVideoBilling.js

@@ -0,0 +1,210 @@
+"use strict";
+
+const ENTERPRISE_VIDEO_ALLOWED_MODELS = new Set([
+  "happyhorse-1.0",
+  "happyhorse-1.0-t2v",
+  "happyhorse-1.0-i2v",
+  "happyhorse-1.0-r2v",
+  "wan2.7-i2v",
+  "wan2.2-animate-mix",
+  "wan2.2-s2v",
+  "kling-3.0-dashscope",
+  "kling-v3-omni-dashscope",
+  "kling/kling-v3-omni-video-generation",
+]);
+
+function normalizeModel(value) {
+  return String(value || "").trim().toLowerCase();
+}
+
+function normalizeEnterpriseVideoResolution(value) {
+  return String(value || "").trim().toUpperCase() === "720P" ? "720P" : "1080P";
+}
+
+function normalizeEnterpriseVideoDuration(value) {
+  const numeric = Number(value);
+  if (!Number.isFinite(numeric)) return 1;
+  return Math.max(1, Math.ceil(numeric));
+}
+
+function isEnterpriseVideoBillingUser(user) {
+  return Boolean(user?.enterpriseId);
+}
+
+function isEnterpriseVideoModelAllowed(providerConfig, model) {
+  const requestedModel = normalizeModel(model || providerConfig?.requestedModel || providerConfig?.model);
+  const resolvedModel = normalizeModel(providerConfig?.model);
+  const protocol = String(providerConfig?.protocol || "").toLowerCase();
+
+  if (ENTERPRISE_VIDEO_ALLOWED_MODELS.has(requestedModel)) return true;
+  if (ENTERPRISE_VIDEO_ALLOWED_MODELS.has(resolvedModel)) return true;
+  if (protocol.startsWith("happyhorse-")) return true;
+  if (protocol === "wan-i2v") return true;
+  if (protocol === "wan-animate-mix") return true;
+  if (protocol === "wan-s2v") return true;
+  if (protocol === "kling-dashscope") return true;
+  return false;
+}
+
+function assertEnterpriseVideoModelAllowed(providerConfig, model) {
+  if (isEnterpriseVideoModelAllowed(providerConfig, model)) return;
+  const error = new Error("该企业账号不可使用当前视频模型");
+  error.status = 403;
+  error.code = "ENTERPRISE_VIDEO_MODEL_NOT_ALLOWED";
+  throw error;
+}
+
+function getEnterpriseVideoCreditRate(input) {
+  const resolution = normalizeEnterpriseVideoResolution(input.resolution || input.quality);
+  const model = normalizeModel(input.model || input.requestedModel);
+
+  if (model.includes("happyhorse")) {
+    return resolution === "720P" ? 0.72 : 1.28;
+  }
+
+  if (model.includes("wan2.7-i2v") || model.includes("wanxiang")) {
+    return resolution === "720P" ? 0.6 : 1;
+  }
+
+  if (model.includes("animate-mix") || model.includes("s2v")) {
+    return resolution === "720P" ? 0.6 : 1;
+  }
+
+  if (model.includes("kling")) {
+    if (input.muted) {
+      if (input.hasReferenceVideo) return resolution === "720P" ? 0.9 : 1.2;
+      return resolution === "720P" ? 0.6 : 0.8;
+    }
+    return resolution === "720P" ? 0.9 : 1.2;
+  }
+
+  const error = new Error(`Unsupported enterprise video model: ${input.model || input.requestedModel}`);
+  error.status = 403;
+  error.code = "ENTERPRISE_VIDEO_MODEL_NOT_ALLOWED";
+  throw error;
+}
+
+function calculateEnterpriseVideoCredits(input) {
+  const duration = normalizeEnterpriseVideoDuration(input.durationSeconds || input.duration);
+  return Number((getEnterpriseVideoCreditRate(input) * duration).toFixed(2));
+}
+
+function calculateEnterpriseVideoCost(input) {
+  const resolution = normalizeEnterpriseVideoResolution(input.resolution || input.quality);
+  const durationSeconds = normalizeEnterpriseVideoDuration(input.durationSeconds || input.duration);
+  const rateCreditsPerSecond = getEnterpriseVideoCreditRate({
+    ...input,
+    resolution,
+    durationSeconds,
+  });
+  const rateCentsPerSecond = Math.round(rateCreditsPerSecond * 100);
+  return {
+    resolution,
+    durationSeconds,
+    rateCentsPerSecond,
+    amountCents: rateCentsPerSecond * durationSeconds,
+  };
+}
+
+function prepareEnterpriseVideoBilling({ user, providerConfig, params }) {
+  if (!isEnterpriseVideoBillingUser(user)) return null;
+  assertEnterpriseVideoModelAllowed(providerConfig, params.requestedModel || params.model);
+  const hasReferenceVideo = Boolean(params.hasReferenceVideo);
+  const muted = Boolean(params.muted);
+  const pricing = calculateEnterpriseVideoCost({
+    model: params.model,
+    requestedModel: params.requestedModel,
+    resolution: params.resolution || params.quality,
+    duration: params.duration,
+    muted,
+    hasReferenceVideo,
+  });
+
+  return {
+    enterpriseId: user.enterpriseId,
+    userId: user.id,
+    model: params.model,
+    requestedModel: params.requestedModel,
+    muted,
+    hasReferenceVideo,
+    ...pricing,
+  };
+}
+
+async function reserveEnterpriseVideoCredits(client, billing) {
+  if (!billing || billing.amountCents <= 0) return null;
+
+  const { rows: balanceRows } = await client.query(
+    "UPDATE enterprises SET balance_cents = balance_cents - $1, updated_at = NOW() WHERE id = $2 AND enabled = 1 AND balance_cents >= $1 RETURNING balance_cents",
+    [billing.amountCents, billing.enterpriseId],
+  );
+  if (balanceRows.length === 0) {
+    const error = new Error(
+      `企业积分不足，预计需要 ${Number((billing.amountCents / 100).toFixed(2))} 积分`,
+    );
+    error.status = 402;
+    error.code = "INSUFFICIENT_ENTERPRISE_BALANCE";
+    throw error;
+  }
+
+  const {
+    rows: [ledger],
+  } = await client.query(
+    `
+    INSERT INTO credit_ledger (
+      enterprise_id,
+      user_id,
+      task_id,
+      model,
+      task_type,
+      resolution,
+      duration_seconds,
+      rate_cents_per_second,
+      amount_cents,
+      status
+    )
+    VALUES ($1, $2, $3, $4, 'video', $5, $6, $7, $8, 'reserved')
+    RETURNING id
+  `,
+    [
+      billing.enterpriseId,
+      billing.userId,
+      billing.taskId || null,
+      billing.requestedModel || billing.model,
+      billing.resolution,
+      billing.durationSeconds,
+      billing.rateCentsPerSecond,
+      billing.amountCents,
+    ],
+  );
+
+  return {
+    ...billing,
+    creditLedgerId: ledger.id,
+    enterpriseBalanceCents: Number(balanceRows[0].balance_cents || 0),
+  };
+}
+
+async function markEnterpriseVideoCreditsAccepted(clientOrPool, creditLedgerId) {
+  if (!creditLedgerId) return false;
+  const { rowCount } = await clientOrPool.query(
+    "UPDATE credit_ledger SET status = 'charged', updated_at = NOW() WHERE id = $1 AND status = 'reserved'",
+    [creditLedgerId],
+  );
+  return rowCount > 0;
+}
+
+module.exports = {
+  ENTERPRISE_VIDEO_ALLOWED_MODELS,
+  assertEnterpriseVideoModelAllowed,
+  calculateEnterpriseVideoCost,
+  calculateEnterpriseVideoCredits,
+  getEnterpriseVideoCreditRate,
+  isEnterpriseVideoBillingUser,
+  isEnterpriseVideoModelAllowed,
+  markEnterpriseVideoCreditsAccepted,
+  normalizeEnterpriseVideoDuration,
+  normalizeEnterpriseVideoResolution,
+  prepareEnterpriseVideoBilling,
+  reserveEnterpriseVideoCredits,
+};