Build a Safe, Metered AI Image and Video Generator with Creem

pnpm dlx degit armitage-labs/creem/packages/examples/safe-metered-ai-generator safe-metered-ai-generator
cd safe-metered-ai-generator
pnpm install

cp .env.example .env.local

pnpm dev

import { Creem } from "creem";
import { env } from "../env";

export const creemClient = new Creem({
  apiKey: env.creemApiKey,
  server: env.isTestMode ? "test" : "prod",
});

-- maps one app user to one Creem credit wallet.
-- account_id is a Creem cca_… id; provider is always 'creem'.
CREATE TABLE IF NOT EXISTS "credit_account" (
  "user_id"           TEXT PRIMARY KEY REFERENCES "user"("id") ON DELETE CASCADE,
  "creem_customer_id" TEXT NOT NULL,
  "account_id"        TEXT NOT NULL,
  "provider"          TEXT NOT NULL DEFAULT 'creem',
  "unit_label"        TEXT NOT NULL DEFAULT 'credits',
  "status"            TEXT NOT NULL DEFAULT 'active',
  -- ...
);

-- a denormalised mirror of every wallet movement, for fast history and audit.
CREATE TABLE IF NOT EXISTS "credit_ledger_cache" (
  "id"             TEXT PRIMARY KEY,
  "user_id"        TEXT NOT NULL REFERENCES "user"("id") ON DELETE CASCADE,
  "account_id"     TEXT NOT NULL,
  "transaction_id" TEXT,
  "side"           TEXT NOT NULL,   -- 'credit' | 'debit'
  "amount"         NUMERIC NOT NULL,
  "reference"      TEXT,
  "kind"           TEXT,            -- pack_purchase | generation | refund | adjustment
  -- ...
);

pnpm db:migrate

export async function screenPrompt(prompt: string, externalId?: string): Promise<ModerationResult> {
  try {
    const data = await creemClient.moderation.screenPrompt(
      { prompt, externalId },
      { timeoutMs: 5000 }, // a slow moderator must not hold the request open forever
    );
    const decision = data.decision;

    if (decision === "allow") {
      return { decision: "allow", allowed: true, raw: data };
    }
    if (decision === "deny" || decision === "flag") {
      return blocked(decision, data);
    }
    // Unknown or absent decision. Do not assume safe. Fail closed.
    return blocked("error", data);
  } catch {
    // API error, network error, or timeout. Fail closed.
    return blocked("error");
  }
}

export async function creditPack(args: {
  userId: string;
  accountId: string;
  amount: number;
  reference: string;
  idempotencyKey: string;
}): Promise<Transaction> {
  const result = await credits.credit(args.accountId, {
    amount: String(args.amount),
    reference: args.reference,
    idempotencyKey: args.idempotencyKey,
  });
  await cacheMovement({
    /* mirror into credit_ledger_cache */
  });
  return result;
}

async debit(accountId, input) {
  try {
    const txn = await creemClient.customerCredits.debitAccount(accountId, input)
    return toTransaction(txn)
  } catch (err) {
    if (isInsufficient(err)) throw new InsufficientBalanceError()
    throw err
  }
}

async getBalance(accountId: string, at?: string): Promise<Balance> {
  const balance = await creemClient.customerCredits.getAccountBalance(accountId, at)
  return toBalance(balance)
}

export async function refundDebit(args: {
  userId: string;
  accountId: string;
  transactionId: string;
  amount: number;
  reference: string;
}): Promise<Transaction> {
  const result = await credits.reverse(args.accountId, args.transactionId);
  await cacheMovement({ ...args, side: "credit", kind: "refund" });
  return result;
}

if (imageDataUrl) {
  const bytes = dataUrlByteLength(imageDataUrl);
  if (bytes === null)
    return NextResponse.json(
      { error: "invalid_image", message: "Reference image must be a base64 data URL." },
      { status: 400 },
    );
  if (bytes > MAX_IMAGE_BYTES)
    return NextResponse.json(
      { error: "image_too_large", message: `Reference image must be under ${MAX_IMAGE_LABEL}.` },
      { status: 413 },
    );
}

// Idempotency: if this key already produced a completed generation, replay it.
const prior = await queryOne(
  `SELECT id, status, result_url FROM generation
   WHERE user_id=$1 AND idempotency_key=$2 ORDER BY created_at DESC LIMIT 1`,
  [user.id, idempotencyKey],
);
if (prior && prior.status === "completed") {
  return NextResponse.json({
    id: prior.id,
    status: "completed",
    url: prior.result_url,
    replayed: true,
  });
}

const moderation = await screenPrompt(prompt, `user_${user.id}:gen_${genId}`);
if (!moderation.allowed) {
  await recordGeneration({
    /* status: 'rejected', decision: moderation.decision */
  });
  return NextResponse.json(
    { id: genId, status: "rejected", decision: moderation.decision, message: moderation.reason },
    { status: 400 },
  );
}

let debitTxnId: string;
try {
  const txn = await debitForGeneration({
    userId: user.id,
    accountId: account.accountId,
    amount: cost,
    reference: `gen:${genId}`,
    idempotencyKey: `debit_${idempotencyKey}`,
  });
  debitTxnId = txn.id;
} catch (err) {
  if (err instanceof InsufficientBalanceError) {
    await recordGeneration({
      /* status: 'insufficient_credits' */
    });
    return NextResponse.json(
      {
        status: "insufficient_credits",
        message: `You need ${cost} credits for this ${mediaType}.`,
      },
      { status: 402 },
    );
  }
  return NextResponse.json({ error: "billing_error" }, { status: 502 });
}

try {
  const generator = getGenerator();
  const result = await generator.generate({
    prompt,
    mediaType,
    imageDataUrl,
    options,
    requestId: genId,
  });
  // 4. RETURN
  await query(`UPDATE generation SET status='completed', result_url=$2 WHERE id=$1`, [
    genId,
    result.url,
  ]);
  return NextResponse.json({ id: genId, status: "completed", url: result.url, cost });
} catch (err) {
  // Model failed AFTER we charged. Reverse the debit, but only claim a refund
  // once the reversal actually lands.
  try {
    await refundDebit({
      userId: user.id,
      accountId: account.accountId,
      transactionId: debitTxnId,
      amount: cost,
      reference: `refund:${genId}`,
    });
  } catch (refundErr) {
    // Reversal failed: the user is still charged. Do not claim a refund that
    // never happened. Mark it failed for reconciliation and say so.
    await query(`UPDATE generation SET status='failed', error=$2 WHERE id=$1`, [
      genId,
      "generation+refund failed",
    ]);
    return NextResponse.json(
      {
        status: "failed",
        message:
          "Generation failed and your credits could not be automatically refunded - our team will reconcile this shortly.",
      },
      { status: 502 },
    );
  }
  // Reversal succeeded. Now it is safe to record and report the refund.
  await query(`UPDATE generation SET status='refunded', error=$2 WHERE id=$1`, [
    genId,
    String(err.message).slice(0, 500),
  ]);
  return NextResponse.json(
    { status: "refunded", message: "Generation failed - your credits were refunded." },
    { status: 502 },
  );
}

const lockName = `gen:${user.id}:${idempotencyKey}`;
const lockClient = await pool.connect();
try {
  await lockClient.query("SELECT pg_advisory_lock(hashtext($1))", [lockName]);
  return await runGeneration();
} finally {
  try {
    await lockClient.query("SELECT pg_advisory_unlock(hashtext($1))", [lockName]);
    lockClient.release();
  } catch (unlockErr) {
    // Destroy the connection so a session lock that failed to release can never
    // return to the pool still held.
    lockClient.release(unlockErr as Error);
  }
}

export const PACKS: Pack[] = [
  { id: "starter", name: "Starter", credits: 200, priceUsd: 9, productId: env.products.starter },
  {
    id: "pro",
    name: "Pro",
    credits: 1000,
    priceUsd: 39,
    productId: env.products.pro,
    featured: true,
  },
  { id: "studio", name: "Studio", credits: 3000, priceUsd: 99, productId: env.products.studio },
];

const { url } = await createCheckout(
  { apiKey: env.creemApiKey, testMode: env.isTestMode },
  {
    productId: pack.productId,
    customer: { email: user.email },
    successUrl: `${env.appUrl}/success`,
    ...(DISCOUNT_CODE ? { discountCode: DISCOUNT_CODE } : {}),
    metadata: { referenceId: user.id, packId: pack.id },
  },
);

onCheckoutCompleted: async ({ customer, product, order, metadata }) => {
  const customerId = customer?.id;
  const userId = metadata?.referenceId as string | undefined;
  if (!customerId || !userId) return;

  const pack = getPackByProductId(product?.id ?? "");
  if (!pack) return;

  // Ensure the wallet exists for the now-known Creem customer, then credit it.
  const account = await ensureAccountForUser(userId, customerId);
  await creditPack({
    userId,
    accountId: account.accountId,
    amount: pack.credits,
    reference: `pack:${pack.id}:order:${order?.id ?? "unknown"}`,
    idempotencyKey: `pack_${order?.id ?? `${userId}_${pack.id}`}`,
  });
};

try {
  if (parsed.data.action === "freeze") await freezeAccount(user.id);
  else if (parsed.data.action === "unfreeze") await unfreezeAccount(user.id);
  else await closeAccount(user.id);
} catch (err) {
  console.error("[account] lifecycle action failed:", err);
  return NextResponse.json(
    { error: "action_failed", message: (err as Error).message },
    { status: 502 },
  );
}

// src/lib/generator/fal.ts
export class FalGenerator implements Generator {
  readonly name = "fal";
  async generate({
    prompt,
    mediaType,
    imageDataUrl,
    options,
  }: GenerateInput): Promise<GenerateResult> {
    // call fal, return { url, mediaType }
  }
}