ai-chatbot/lib/video-agent/task-store.ts
Your Name ec988be30a
Some checks failed
Lint / build (20) (push) Has been cancelled
Playwright Tests / test (push) Has been cancelled
Update agent workflow and video task handling
2026-03-13 11:21:42 +08:00

720 lines
19 KiB
TypeScript

import "server-only";
import {
and,
asc,
desc,
eq,
inArray,
sql,
type SQL,
} from "drizzle-orm";
import { db } from "@/lib/db";
import {
videoStoryboardShotTask,
videoStoryboardTask,
} from "@/lib/db/schema";
import type { PreparedStoryboardShot } from "@/lib/video-agent/types";
export type VideoStoryboardTaskStatus =
| "pending"
| "running"
| "partial_success"
| "succeeded"
| "failed"
| "cancelled";
export type VideoStoryboardTaskConfirmationStatus =
| "draft"
| "awaiting_user_confirmation"
| "confirmed"
| "rejected";
export type VideoStoryboardShotStatus =
| "pending"
| "running"
| "completed"
| "failed"
| "cancelled";
export type VideoStoryboardHdImageStatus =
| "pending"
| "running"
| "completed"
| "failed"
| "cancelled";
export type VideoStoryboardVideoStatus =
| "pending"
| "waiting_confirmation"
| "running"
| "completed"
| "failed"
| "cancelled";
function toConfidenceInt(confidence: number) {
if (!Number.isFinite(confidence)) return 0;
return Math.max(0, Math.min(10000, Math.round(confidence * 10000)));
}
export async function createVideoStoryboardTask(params: {
userId: string;
chatId?: string;
providerType: string;
modelRef?: string;
gridType?: string | null;
gridRows?: number | null;
gridCols?: number | null;
sourceImageUrl: string;
isNineGrid: boolean;
detectionConfidence: number;
normalizedScript?: Record<string, any> | null;
totalShots: number;
status?: VideoStoryboardTaskStatus;
confirmationStatus?: VideoStoryboardTaskConfirmationStatus;
}) {
const now = new Date();
const [task] = await db
.insert(videoStoryboardTask)
.values({
userId: params.userId,
chatId: params.chatId || null,
providerType: params.providerType,
modelRef: params.modelRef || null,
gridType: params.gridType || null,
gridRows:
typeof params.gridRows === "number" && Number.isFinite(params.gridRows)
? Math.max(1, Math.round(params.gridRows))
: null,
gridCols:
typeof params.gridCols === "number" && Number.isFinite(params.gridCols)
? Math.max(1, Math.round(params.gridCols))
: null,
sourceImageUrl: params.sourceImageUrl,
isNineGrid: params.isNineGrid,
detectionConfidence: toConfidenceInt(params.detectionConfidence),
normalizedScript: params.normalizedScript || null,
status: params.status || "running",
confirmationStatus: params.confirmationStatus || "draft",
hdImageStageStatus: "pending",
currentShot: 0,
totalShots: Math.max(0, Math.round(params.totalShots)),
createdAt: now,
updatedAt: now,
startedAt: now,
})
.returning();
return task;
}
export async function createVideoStoryboardShotTasks(params: {
parentTaskId: string;
shots: PreparedStoryboardShot[];
fallbackChain: string[];
}) {
const now = new Date();
if (!params.shots.length) return [];
const inserted = await db
.insert(videoStoryboardShotTask)
.values(
params.shots.map((shot) => ({
parentTaskId: params.parentTaskId,
shotIndex: shot.shotIndex,
tileIndex: shot.tileIndex,
tileImageUrl: shot.tileImageUrl,
upscaledImageUrl: shot.upscaled ? shot.inputImageUrl : null,
hdImageUrl: null,
hdImagePrompt: shot.prompt,
prompt: shot.prompt,
modelId: null,
fallbackChain: params.fallbackChain,
attemptCount: 0,
status: "pending",
hdImageStatus: "pending",
videoStatus: "pending",
userApproved: false,
userApprovedAt: null,
createdAt: now,
updatedAt: now,
}))
)
.returning();
return inserted;
}
export async function attachVideoStoryboardTaskMessage(params: {
taskId: string;
messageId: string;
}) {
await db
.update(videoStoryboardTask)
.set({
messageId: params.messageId,
updatedAt: new Date(),
})
.where(eq(videoStoryboardTask.id, params.taskId));
}
export async function markVideoStoryboardShotRunning(params: {
taskId: string;
shotIndex: number;
modelId: string;
}) {
const now = new Date();
await db
.update(videoStoryboardShotTask)
.set({
status: "running",
modelId: params.modelId,
startedAt: now,
completedAt: null,
errorCode: null,
errorMessage: null,
updatedAt: now,
attemptCount: sql`${videoStoryboardShotTask.attemptCount} + 1`,
})
.where(
and(
eq(videoStoryboardShotTask.parentTaskId, params.taskId),
eq(videoStoryboardShotTask.shotIndex, params.shotIndex)
)
);
await db
.update(videoStoryboardTask)
.set({
status: "running",
currentShot: params.shotIndex,
updatedAt: now,
})
.where(eq(videoStoryboardTask.id, params.taskId));
}
export async function markVideoStoryboardShotRetry(params: {
taskId: string;
shotIndex: number;
errorMessage: string;
}) {
await db
.update(videoStoryboardShotTask)
.set({
status: "running",
errorMessage: params.errorMessage.slice(0, 4000),
updatedAt: new Date(),
})
.where(
and(
eq(videoStoryboardShotTask.parentTaskId, params.taskId),
eq(videoStoryboardShotTask.shotIndex, params.shotIndex)
)
);
}
export async function markVideoStoryboardShotCompleted(params: {
taskId: string;
shotIndex: number;
modelId: string;
videoUrl: string;
}) {
const now = new Date();
await db
.update(videoStoryboardShotTask)
.set({
status: "completed",
modelId: params.modelId,
videoUrl: params.videoUrl,
errorCode: null,
errorMessage: null,
completedAt: now,
updatedAt: now,
})
.where(
and(
eq(videoStoryboardShotTask.parentTaskId, params.taskId),
eq(videoStoryboardShotTask.shotIndex, params.shotIndex)
)
);
}
export async function markVideoStoryboardShotFailed(params: {
taskId: string;
shotIndex: number;
errorMessage: string;
}) {
const now = new Date();
await db
.update(videoStoryboardShotTask)
.set({
status: "failed",
errorCode: "shot_failed",
errorMessage: params.errorMessage.slice(0, 4000),
completedAt: now,
updatedAt: now,
})
.where(
and(
eq(videoStoryboardShotTask.parentTaskId, params.taskId),
eq(videoStoryboardShotTask.shotIndex, params.shotIndex)
)
);
}
export async function finalizeVideoStoryboardTask(params: {
taskId: string;
status: VideoStoryboardTaskStatus;
error?: string;
}) {
const now = new Date();
await db
.update(videoStoryboardTask)
.set({
status: params.status,
error: params.error ? params.error.slice(0, 4000) : null,
completedAt:
params.status === "succeeded" ||
params.status === "partial_success" ||
params.status === "failed" ||
params.status === "cancelled"
? now
: null,
updatedAt: now,
})
.where(eq(videoStoryboardTask.id, params.taskId));
}
export async function markVideoStoryboardAwaitingConfirmation(params: {
taskId: string;
}) {
const now = new Date();
await db
.update(videoStoryboardTask)
.set({
confirmationStatus: "awaiting_user_confirmation",
status: "pending",
updatedAt: now,
})
.where(eq(videoStoryboardTask.id, params.taskId));
await db
.update(videoStoryboardShotTask)
.set({
videoStatus: "waiting_confirmation",
userApproved: false,
userApprovedAt: null,
updatedAt: now,
})
.where(eq(videoStoryboardShotTask.parentTaskId, params.taskId));
}
export async function markVideoStoryboardHdRunning(params: {
taskId: string;
shotIndex: number;
}) {
const now = new Date();
await db
.update(videoStoryboardShotTask)
.set({
hdImageStatus: "running",
updatedAt: now,
})
.where(
and(
eq(videoStoryboardShotTask.parentTaskId, params.taskId),
eq(videoStoryboardShotTask.shotIndex, params.shotIndex)
)
);
await db
.update(videoStoryboardTask)
.set({
hdImageStageStatus: "running",
updatedAt: now,
})
.where(eq(videoStoryboardTask.id, params.taskId));
}
export async function markVideoStoryboardHdCompleted(params: {
taskId: string;
shotIndex: number;
hdImageUrl: string;
hdImagePrompt?: string;
}) {
const now = new Date();
await db
.update(videoStoryboardShotTask)
.set({
hdImageStatus: "completed",
hdImageUrl: params.hdImageUrl,
hdImagePrompt: params.hdImagePrompt || null,
updatedAt: now,
})
.where(
and(
eq(videoStoryboardShotTask.parentTaskId, params.taskId),
eq(videoStoryboardShotTask.shotIndex, params.shotIndex)
)
);
}
export async function markVideoStoryboardHdFailed(params: {
taskId: string;
shotIndex: number;
errorMessage: string;
}) {
const now = new Date();
await db
.update(videoStoryboardShotTask)
.set({
hdImageStatus: "failed",
errorMessage: params.errorMessage.slice(0, 4000),
updatedAt: now,
})
.where(
and(
eq(videoStoryboardShotTask.parentTaskId, params.taskId),
eq(videoStoryboardShotTask.shotIndex, params.shotIndex)
)
);
}
export async function syncVideoStoryboardHdStageStatus(taskId: string) {
const rows = await db
.select({
hdImageStatus: videoStoryboardShotTask.hdImageStatus,
})
.from(videoStoryboardShotTask)
.where(eq(videoStoryboardShotTask.parentTaskId, taskId));
const total = rows.length;
const completed = rows.filter((row) => row.hdImageStatus === "completed").length;
const failed = rows.filter((row) => row.hdImageStatus === "failed").length;
const running = rows.filter((row) => row.hdImageStatus === "running").length;
let hdImageStageStatus: "pending" | "running" | "completed" | "failed" = "pending";
if (running > 0) {
hdImageStageStatus = "running";
} else if (completed === total && total > 0) {
hdImageStageStatus = "completed";
} else if (failed === total && total > 0) {
hdImageStageStatus = "failed";
} else if (completed > 0 || failed > 0) {
hdImageStageStatus = "running";
}
await db
.update(videoStoryboardTask)
.set({
hdImageStageStatus,
updatedAt: new Date(),
})
.where(eq(videoStoryboardTask.id, taskId));
return { total, completed, failed, hdImageStageStatus };
}
export async function markVideoStoryboardConfirmed(params: {
taskId: string;
messageId?: string;
}) {
const now = new Date();
await db
.update(videoStoryboardTask)
.set({
confirmationStatus: "confirmed",
confirmedAt: now,
confirmedByMessageId: params.messageId || null,
updatedAt: now,
})
.where(eq(videoStoryboardTask.id, params.taskId));
await db
.update(videoStoryboardShotTask)
.set({
videoStatus: sql`CASE
WHEN ${videoStoryboardShotTask.status} = 'completed' THEN 'completed'
WHEN ${videoStoryboardShotTask.status} = 'failed' THEN 'failed'
WHEN ${videoStoryboardShotTask.status} = 'running' THEN 'running'
ELSE 'pending'
END`,
userApproved: true,
userApprovedAt: now,
updatedAt: now,
})
.where(eq(videoStoryboardShotTask.parentTaskId, params.taskId));
}
export async function markVideoStoryboardDraftForRevision(params: {
taskId: string;
}) {
const now = new Date();
await db
.update(videoStoryboardTask)
.set({
confirmationStatus: "draft",
confirmedAt: null,
confirmedByMessageId: null,
updatedAt: now,
})
.where(eq(videoStoryboardTask.id, params.taskId));
await db
.update(videoStoryboardShotTask)
.set({
videoStatus: "waiting_confirmation",
userApproved: false,
userApprovedAt: null,
updatedAt: now,
})
.where(eq(videoStoryboardShotTask.parentTaskId, params.taskId));
}
export async function updateVideoStoryboardShotPrompt(params: {
taskId: string;
shotIndex: number;
prompt: string;
}) {
const now = new Date();
await db
.update(videoStoryboardShotTask)
.set({
prompt: params.prompt.slice(0, 4000),
hdImagePrompt: params.prompt.slice(0, 4000),
updatedAt: now,
})
.where(
and(
eq(videoStoryboardShotTask.parentTaskId, params.taskId),
eq(videoStoryboardShotTask.shotIndex, params.shotIndex)
)
);
}
export async function cancelVideoStoryboardTaskByUser(params: {
taskId: string;
error?: string;
}) {
const now = new Date();
await db
.update(videoStoryboardTask)
.set({
status: "cancelled",
confirmationStatus: "rejected",
error: params.error ? params.error.slice(0, 4000) : "user_cancelled",
updatedAt: now,
completedAt: now,
})
.where(eq(videoStoryboardTask.id, params.taskId));
await db
.update(videoStoryboardShotTask)
.set({
status: sql`CASE
WHEN ${videoStoryboardShotTask.status} IN ('completed', 'failed') THEN ${videoStoryboardShotTask.status}
ELSE 'cancelled'
END`,
videoStatus: sql`CASE
WHEN ${videoStoryboardShotTask.videoStatus} IN ('completed', 'failed') THEN ${videoStoryboardShotTask.videoStatus}
ELSE 'cancelled'
END`,
userApproved: false,
userApprovedAt: null,
updatedAt: now,
completedAt: sql`CASE
WHEN ${videoStoryboardShotTask.completedAt} IS NULL THEN ${now}
ELSE ${videoStoryboardShotTask.completedAt}
END`,
})
.where(eq(videoStoryboardShotTask.parentTaskId, params.taskId));
}
export async function syncVideoStoryboardTaskStatus(taskId: string) {
const rows = await db
.select({
status: videoStoryboardShotTask.status,
})
.from(videoStoryboardShotTask)
.where(eq(videoStoryboardShotTask.parentTaskId, taskId));
const total = rows.length;
const success = rows.filter((row) => row.status === "completed").length;
const failed = rows.filter((row) => row.status === "failed").length;
const running = rows.filter((row) => row.status === "running").length;
const pending = rows.filter((row) => row.status === "pending").length;
let status: VideoStoryboardTaskStatus = "pending";
if (running > 0) {
status = "running";
} else if (success === total && total > 0) {
status = "succeeded";
} else if (failed === total && total > 0) {
status = "failed";
} else if (success > 0 && failed > 0) {
status = "partial_success";
} else if (success > 0 && pending > 0) {
status = "running";
} else if (failed > 0 && pending > 0) {
status = "running";
}
const currentShot = await db
.select({
maxShot: sql<number>`max(${videoStoryboardShotTask.shotIndex})::int`,
})
.from(videoStoryboardShotTask)
.where(
and(
eq(videoStoryboardShotTask.parentTaskId, taskId),
inArray(videoStoryboardShotTask.status, ["running", "completed", "failed"])
)
);
await db
.update(videoStoryboardTask)
.set({
status,
currentShot: Number(currentShot[0]?.maxShot || 0),
completedAt:
status === "succeeded" || status === "partial_success" || status === "failed"
? new Date()
: null,
updatedAt: new Date(),
})
.where(eq(videoStoryboardTask.id, taskId));
return { total, success, failed, status };
}
export async function getVideoStoryboardTaskDetail(taskId: string) {
const taskRows = await db
.select()
.from(videoStoryboardTask)
.where(eq(videoStoryboardTask.id, taskId))
.limit(1);
const task = taskRows[0];
if (!task) return null;
const shots = await db
.select()
.from(videoStoryboardShotTask)
.where(eq(videoStoryboardShotTask.parentTaskId, taskId))
.orderBy(asc(videoStoryboardShotTask.shotIndex));
return { task, shots };
}
export async function getLatestActiveVideoStoryboardTaskByChatId(params: {
chatId: string;
}) {
const taskRows = await db
.select()
.from(videoStoryboardTask)
.where(
and(
eq(videoStoryboardTask.chatId, params.chatId),
inArray(videoStoryboardTask.status, ["pending", "running"])
)
)
.orderBy(desc(videoStoryboardTask.createdAt))
.limit(1);
const task = taskRows[0];
if (!task) return null;
const shots = await db
.select()
.from(videoStoryboardShotTask)
.where(eq(videoStoryboardShotTask.parentTaskId, task.id))
.orderBy(asc(videoStoryboardShotTask.shotIndex));
return { task, shots };
}
export async function listVideoStoryboardTasks(params: {
page: number;
pageSize: number;
status?: string;
providerType?: string;
chatId?: string;
userId?: string;
}) {
const conditions: SQL[] = [];
if (params.status) {
conditions.push(eq(videoStoryboardTask.status, params.status));
}
if (params.providerType) {
conditions.push(eq(videoStoryboardTask.providerType, params.providerType));
}
if (params.chatId) {
conditions.push(eq(videoStoryboardTask.chatId, params.chatId));
}
if (params.userId) {
conditions.push(eq(videoStoryboardTask.userId, params.userId));
}
const whereClause = conditions.length > 0 ? and(...conditions) : undefined;
const page = Math.max(1, Math.round(params.page || 1));
const pageSize = Math.max(1, Math.min(200, Math.round(params.pageSize || 20)));
const offset = (page - 1) * pageSize;
const [countRow] = await db
.select({ count: sql<number>`count(*)::int` })
.from(videoStoryboardTask)
.where(whereClause);
const tasks = await db
.select()
.from(videoStoryboardTask)
.where(whereClause)
.orderBy(desc(videoStoryboardTask.createdAt))
.limit(pageSize)
.offset(offset);
if (tasks.length === 0) {
return {
total: Number(countRow?.count || 0),
page,
pageSize,
tasks: [] as Array<any>,
};
}
const taskIds = tasks.map((item) => item.id);
const shotRows = await db
.select({
parentTaskId: videoStoryboardShotTask.parentTaskId,
status: videoStoryboardShotTask.status,
})
.from(videoStoryboardShotTask)
.where(inArray(videoStoryboardShotTask.parentTaskId, taskIds));
const summaryMap = new Map<
string,
{ totalShots: number; completedShots: number; failedShots: number; runningShots: number }
>();
for (const row of shotRows) {
const summary =
summaryMap.get(row.parentTaskId) ||
{ totalShots: 0, completedShots: 0, failedShots: 0, runningShots: 0 };
summary.totalShots += 1;
if (row.status === "completed") summary.completedShots += 1;
if (row.status === "failed") summary.failedShots += 1;
if (row.status === "running") summary.runningShots += 1;
summaryMap.set(row.parentTaskId, summary);
}
return {
total: Number(countRow?.count || 0),
page,
pageSize,
tasks: tasks.map((task) => ({
...task,
summary: summaryMap.get(task.id) || {
totalShots: 0,
completedShots: 0,
failedShots: 0,
runningShots: 0,
},
})),
};
}