douyin-archive/lib/stt/tasks.ts

import { randomUUID } from "node:crypto";
import { Prisma, type TranscriptionTask } from "@prisma/client";
import { prisma } from "../prisma";
import { getAsrConfig } from "./dashscope";
import { isActive, transcriptionState } from "./types";

function configurationError(): string | null {
  try {
    getAsrConfig();
    return null;
  } catch (error) {
    return (error as Error).message;
  }
}
export async function enqueueTranscription(videoId: string, force = false) {
  return prisma.$transaction(async (tx) => {
    const videos = await tx.$queryRaw<
      { aweme_id: string }[]
    >`SELECT aweme_id FROM "Video" WHERE aweme_id = ${videoId} FOR UPDATE`;
    if (!videos.length) return null;
    const [transcript, task] = await Promise.all([
      tx.videoTranscript.findUnique({ where: { videoId } }),
      tx.transcriptionTask.findUnique({ where: { videoId } }),
    ]);
    if ((task && isActive(task.status)) || (!force && transcript))
      return transcriptionState(transcript, task);
    const lastError = configurationError();
    const data = {
      status: lastError ? "FAILED" : "QUEUED",
      lastError,
      providerTaskId: null,
      submittingAt: null,
      submittedAt: null,
      leaseUntil: null,
      leaseToken: null,
      attempts: 0,
      nextPollAt: new Date(),
    };
    const saved = await tx.transcriptionTask.upsert({
      where: { videoId },
      create: { videoId, ...data },
      update: data,
    });
    return transcriptionState(transcript, saved);
  });
}
export async function enqueueMissingTranscriptions() {
  const error = configurationError();
  if (error) throw new Error(error);
  // One durable, atomic batch. Completed historical transcripts are never requeued.
  return prisma.$executeRaw`
    INSERT INTO "TranscriptionTask" ("videoId", status, "updatedAt")
    SELECT v.aweme_id, 'QUEUED', (CURRENT_TIMESTAMP AT TIME ZONE 'UTC') FROM "Video" v
    WHERE NOT EXISTS (SELECT 1 FROM "VideoTranscript" t WHERE t."videoId" = v.aweme_id)
    ON CONFLICT ("videoId") DO UPDATE SET status = 'QUEUED', "providerTaskId" = NULL,
      "submittingAt" = NULL, "submittedAt" = NULL, "leaseUntil" = NULL, "leaseToken" = NULL,
      attempts = 0, "lastError" = NULL, "nextPollAt" = (CURRENT_TIMESTAMP AT TIME ZONE 'UTC'), "updatedAt" = (CURRENT_TIMESTAMP AT TIME ZONE 'UTC')
    WHERE "TranscriptionTask".status = 'FAILED'
  `;
}
export async function readTranscription(videoId: string) {
  const video = await prisma.video.findUnique({
    where: { aweme_id: videoId },
    include: { transcript: true, transcriptionTask: true },
  });
  return video
    ? transcriptionState(video.transcript, video.transcriptionTask)
    : null;
}

export async function claimTask(): Promise<TranscriptionTask | null> {
  return prisma.$transaction(async (tx) => {
    // Serialize admission across all worker processes, including PM2 restarts.
    await tx.$queryRaw`SELECT pg_advisory_xact_lock(74196321)::text`;
    const running = await tx.transcriptionTask.count({
      where: { status: "PROCESSING" },
    });
    const eligible =
      running < 2
        ? Prisma.sql`status IN ('QUEUED', 'PROCESSING')`
        : Prisma.sql`status = 'PROCESSING'`;
    const rows = await tx.$queryRaw<TranscriptionTask[]>`
      SELECT * FROM "TranscriptionTask" WHERE ${eligible}
        AND "nextPollAt" <= (CURRENT_TIMESTAMP AT TIME ZONE 'UTC') AND ("leaseUntil" IS NULL OR "leaseUntil" < (CURRENT_TIMESTAMP AT TIME ZONE 'UTC'))
      ORDER BY CASE WHEN status = 'PROCESSING' THEN 0 ELSE 1 END, "nextPollAt", "createdAt"
      LIMIT 1 FOR UPDATE SKIP LOCKED
    `;
    if (!rows.length) return null;
    return tx.transcriptionTask.update({
      where: { videoId: rows[0].videoId },
      data: {
        status: "PROCESSING",
        leaseToken: randomUUID(),
        leaseUntil: new Date(Date.now() + 180_000),
      },
    });
  });
}
export async function updateOwnedTask(
  task: TranscriptionTask,
  data: Prisma.TranscriptionTaskUpdateManyMutationInput,
) {
  const result = await prisma.transcriptionTask.updateMany({
    where: {
      videoId: task.videoId,
      leaseToken: task.leaseToken,
      status: "PROCESSING",
      leaseUntil: { gt: new Date() },
    },
    data,
  });
  if (!result.count) throw new Error("转写任务租约已失效");
}
export async function completeTask(
  task: TranscriptionTask,
  transcript: string[],
) {
  await prisma.$transaction(async (tx) => {
    const result = await tx.transcriptionTask.updateMany({
      where: {
        videoId: task.videoId,
        leaseToken: task.leaseToken,
        status: "PROCESSING",
        leaseUntil: { gt: new Date() },
      },
      data: {
        status: "SUCCEEDED",
        lastError: null,
        leaseToken: null,
        leaseUntil: null,
        attempts: 0,
      },
    });
    if (!result.count) throw new Error("转写任务租约已失效");
    await tx.videoTranscript.upsert({
      where: { videoId: task.videoId },
      create: { videoId: task.videoId, transcript },
      update: { transcript },
    });
  });
}