351 lines
11 KiB
TypeScript
351 lines
11 KiB
TypeScript
import { db } from "@/lib/db";
|
|
import { generatedImage } from "@/lib/db/schema";
|
|
import {
|
|
consumeUserPoints,
|
|
createMediaAssetRelation,
|
|
getAiModelById,
|
|
getEnabledAiModelsWithProvider,
|
|
getMediaAssetsByUserAndUrls,
|
|
getPointsSystemConfig,
|
|
getUserById,
|
|
upsertMediaAsset,
|
|
} from "@/lib/db/queries";
|
|
import { persistImageWorkflowTurn } from "@/lib/ai/conversation-state-protocol";
|
|
import { getModelGenerationCreditCost } from "@/lib/ai/image-model-pricing";
|
|
|
|
const MODEL_UUID_REGEX =
|
|
/^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$/i;
|
|
|
|
export function parseImageAltMetadata(altText: string, thumbnailUrl: string): {
|
|
originalUrl: string;
|
|
imageSize: string;
|
|
} {
|
|
let originalUrl = thumbnailUrl;
|
|
let imageSize = "2K";
|
|
|
|
if (altText.includes("thumb|||")) {
|
|
const content = altText.substring(altText.indexOf("thumb|||") + 8);
|
|
const parts = content.split("|||");
|
|
originalUrl = parts[0] || thumbnailUrl;
|
|
imageSize = parts[1] || "2K";
|
|
} else if (altText.includes("thumb$$")) {
|
|
const content = altText.substring(altText.indexOf("thumb$$") + 7);
|
|
const parts = content.split("$$");
|
|
originalUrl = parts[0] || thumbnailUrl;
|
|
imageSize = parts[1] || "2K";
|
|
} else if (altText.includes("thumb|")) {
|
|
const content = altText.substring(altText.indexOf("thumb|") + 6);
|
|
const parts = content.split("|");
|
|
originalUrl = parts[0] || thumbnailUrl;
|
|
imageSize = parts[1] || "2K";
|
|
}
|
|
|
|
if (originalUrl === thumbnailUrl) {
|
|
if (thumbnailUrl.includes("gemini-thumb-")) {
|
|
originalUrl = thumbnailUrl.replace("gemini-thumb-", "gemini-orig-");
|
|
} else if (thumbnailUrl.includes("openai-thumb-")) {
|
|
originalUrl = thumbnailUrl.replace("openai-thumb-", "openai-orig-");
|
|
} else if (thumbnailUrl.includes("volcano-thumb-")) {
|
|
originalUrl = thumbnailUrl.replace("volcano-thumb-", "volcano-orig-");
|
|
} else if (thumbnailUrl.includes("bfl-image-thumb-")) {
|
|
originalUrl = thumbnailUrl.replace("bfl-image-thumb-", "bfl-image-orig-");
|
|
} else if (thumbnailUrl.includes("alibaba-image-thumb-")) {
|
|
originalUrl = thumbnailUrl.replace("alibaba-image-thumb-", "alibaba-image-orig-");
|
|
} else if (thumbnailUrl.includes("ruxa-image-thumb-")) {
|
|
originalUrl = thumbnailUrl.replace("ruxa-image-thumb-", "ruxa-image-orig-");
|
|
} else if (thumbnailUrl.includes("stepfun-image-thumb-")) {
|
|
originalUrl = thumbnailUrl.replace("stepfun-image-thumb-", "stepfun-image-orig-");
|
|
} else if (thumbnailUrl.includes("fal-image-thumb-")) {
|
|
originalUrl = thumbnailUrl.replace("fal-image-thumb-", "fal-image-orig-");
|
|
} else if (thumbnailUrl.includes("snapgen-image-thumb-")) {
|
|
originalUrl = thumbnailUrl.replace("snapgen-image-thumb-", "snapgen-image-orig-");
|
|
} else if (thumbnailUrl.includes("bananaimage-thumb-")) {
|
|
originalUrl = thumbnailUrl.replace("bananaimage-thumb-", "bananaimage-orig-");
|
|
}
|
|
}
|
|
|
|
return { originalUrl, imageSize };
|
|
}
|
|
|
|
function getModelCreditCost(
|
|
modelConfig: any,
|
|
target: "image" | "video",
|
|
options?: {
|
|
duration?: number | null;
|
|
imageSize?: string | null;
|
|
resolution?: string | null;
|
|
},
|
|
): number {
|
|
return getModelGenerationCreditCost(modelConfig, target, options);
|
|
}
|
|
|
|
async function resolveBillingModelByReference(modelRef: string): Promise<any | undefined> {
|
|
if (!modelRef) return undefined;
|
|
|
|
if (MODEL_UUID_REGEX.test(modelRef)) {
|
|
return await getAiModelById({ id: modelRef }).catch(() => undefined);
|
|
}
|
|
|
|
const allModels = await getEnabledAiModelsWithProvider(true, true).catch(() => []);
|
|
const matched = allModels.find(
|
|
({ model }: any) => model.id === modelRef || model.modelId === modelRef,
|
|
);
|
|
return matched?.model;
|
|
}
|
|
|
|
export async function saveGeneratedImages(
|
|
imageMarkdown: string,
|
|
userId: string,
|
|
chatId: string | null,
|
|
messageId: string | null,
|
|
prompt: string,
|
|
modelId: string,
|
|
options?: {
|
|
sourceAssetIds?: string[];
|
|
sourceImageUrls?: string[];
|
|
relationType?: "generated_from" | "edited_from";
|
|
sourceType?: string;
|
|
projectId?: string | null;
|
|
userPrompt?: string;
|
|
optimizedPrompt?: string;
|
|
requestedImageCount?: number;
|
|
aspectRatio?: string;
|
|
imageSize?: string;
|
|
skipBilling?: boolean;
|
|
},
|
|
): Promise<{
|
|
savedImageCount: number;
|
|
savedImageIds: string[];
|
|
savedAssetIds: string[];
|
|
savedOriginalUrls: string[];
|
|
}> {
|
|
const imageRegex = /!\[([^\]]*)\]\(([^)]+)\)/g;
|
|
const matches = [...imageMarkdown.matchAll(imageRegex)];
|
|
|
|
let savedImageCount = 0;
|
|
const savedImageIds: string[] = [];
|
|
const savedAssetIds: string[] = [];
|
|
const savedOriginalUrls: string[] = [];
|
|
const savedImageSizes: string[] = [];
|
|
let relationSourceAssetIds = Array.isArray(options?.sourceAssetIds)
|
|
? options.sourceAssetIds.filter((assetId): assetId is string => Boolean(assetId))
|
|
: [];
|
|
|
|
if (
|
|
relationSourceAssetIds.length === 0 &&
|
|
Array.isArray(options?.sourceImageUrls) &&
|
|
options.sourceImageUrls.length > 0
|
|
) {
|
|
try {
|
|
const sourceAssets = await getMediaAssetsByUserAndUrls({
|
|
userId,
|
|
urls: options.sourceImageUrls,
|
|
});
|
|
relationSourceAssetIds = sourceAssets.map((asset) => asset.id);
|
|
} catch (sourceLookupError) {
|
|
console.error("[Gallery] Failed to resolve source assets by URL:", sourceLookupError);
|
|
}
|
|
}
|
|
|
|
for (const match of matches) {
|
|
const altText = match[1] || "";
|
|
const thumbnailUrl = match[2];
|
|
const { originalUrl, imageSize } = parseImageAltMetadata(altText, thumbnailUrl);
|
|
|
|
try {
|
|
const [savedImage] = await db
|
|
.insert(generatedImage)
|
|
.values({
|
|
userId,
|
|
chatId: chatId ?? null,
|
|
messageId: messageId ?? null,
|
|
thumbnailUrl,
|
|
originalUrl,
|
|
imageSize,
|
|
prompt,
|
|
modelId,
|
|
})
|
|
.returning();
|
|
|
|
savedImageCount += 1;
|
|
if (savedImage?.id) {
|
|
savedImageIds.push(savedImage.id);
|
|
}
|
|
if (!savedOriginalUrls.includes(originalUrl)) {
|
|
savedOriginalUrls.push(originalUrl);
|
|
}
|
|
savedImageSizes.push(options?.imageSize || imageSize || "2K");
|
|
|
|
const media = await upsertMediaAsset({
|
|
userId,
|
|
chatId,
|
|
projectId: options?.projectId ?? null,
|
|
mediaType: "image",
|
|
sourceType: options?.sourceType || "generated",
|
|
status: "ready",
|
|
originalUrl,
|
|
previewUrl: originalUrl,
|
|
thumbnailUrl,
|
|
displayName: prompt?.slice(0, 80) || "AI 生成图片",
|
|
prompt,
|
|
modelId,
|
|
sourceMessageId: messageId,
|
|
sourceTable: "GeneratedImage",
|
|
sourceRowId: savedImage?.id || null,
|
|
primaryParentAssetId: relationSourceAssetIds[0] || null,
|
|
metadata: {
|
|
imageSize,
|
|
},
|
|
});
|
|
|
|
if (media?.id && !savedAssetIds.includes(media.id)) {
|
|
savedAssetIds.push(media.id);
|
|
}
|
|
|
|
if (media && relationSourceAssetIds.length > 0) {
|
|
for (const sourceAssetId of relationSourceAssetIds) {
|
|
if (!sourceAssetId) continue;
|
|
await createMediaAssetRelation({
|
|
userId,
|
|
chatId,
|
|
projectId: options?.projectId ?? null,
|
|
fromAssetId: sourceAssetId,
|
|
toAssetId: media.id,
|
|
relationType: options?.relationType || "generated_from",
|
|
sourceMessageId: messageId,
|
|
metadata: {
|
|
modelId,
|
|
prompt,
|
|
sourceImageUrls: options?.sourceImageUrls || [],
|
|
},
|
|
});
|
|
}
|
|
}
|
|
|
|
console.log(`[Gallery] Saved image: ${originalUrl.substring(0, 80)}...`);
|
|
} catch (error) {
|
|
console.error("[Gallery] Failed to save image record:", error);
|
|
}
|
|
}
|
|
|
|
if (savedOriginalUrls.length > 0 && chatId && messageId) {
|
|
try {
|
|
const resultAssetIds = savedAssetIds.length > 0 ? savedAssetIds : null;
|
|
await persistImageWorkflowTurn({
|
|
userId,
|
|
chatId,
|
|
projectId: options?.projectId || null,
|
|
currentIntent: options?.relationType === "edited_from" ? "image_edit" : "image_gen",
|
|
currentFocusAssetIds: resultAssetIds,
|
|
currentFocusImageUrls: savedOriginalUrls,
|
|
lastGeneratedAssetIds: resultAssetIds,
|
|
lastGeneratedImageUrls: savedOriginalUrls,
|
|
lastResolvedTargetAssetIds: resultAssetIds,
|
|
lastResolvedTargetImageUrls: savedOriginalUrls,
|
|
lastMessageId: messageId,
|
|
lastUserPrompt: String(options?.userPrompt || prompt || "").trim() || prompt,
|
|
lastOptimizedPrompt:
|
|
String(
|
|
options?.optimizedPrompt ||
|
|
(options?.relationType !== "edited_from" ? prompt : "") ||
|
|
"",
|
|
).trim() || undefined,
|
|
lastRequestedImageCount: options?.requestedImageCount || savedOriginalUrls.length,
|
|
lastAspectRatio: options?.aspectRatio,
|
|
lastImageSize: options?.imageSize,
|
|
lastSourceImageUrls: options?.sourceImageUrls || null,
|
|
lastResultImageUrls: savedOriginalUrls,
|
|
});
|
|
} catch (memoryError) {
|
|
console.warn("[Gallery] Failed to update working state with generated images:", memoryError);
|
|
}
|
|
}
|
|
|
|
if (savedImageCount <= 0) {
|
|
return {
|
|
savedImageCount,
|
|
savedImageIds,
|
|
savedAssetIds,
|
|
savedOriginalUrls,
|
|
};
|
|
}
|
|
|
|
try {
|
|
if (options?.skipBilling) {
|
|
return {
|
|
savedImageCount,
|
|
savedImageIds,
|
|
savedAssetIds,
|
|
savedOriginalUrls,
|
|
};
|
|
}
|
|
|
|
const pointsConfig = await getPointsSystemConfig();
|
|
if (!pointsConfig.enabled) {
|
|
return {
|
|
savedImageCount,
|
|
savedImageIds,
|
|
savedAssetIds,
|
|
savedOriginalUrls,
|
|
};
|
|
}
|
|
|
|
const userData = await getUserById({ id: userId }).catch(() => undefined);
|
|
if (userData?.isHighVip) {
|
|
console.log(`[Points] Skip image points for VIP user=${userId} model=${modelId}`);
|
|
return {
|
|
savedImageCount,
|
|
savedImageIds,
|
|
savedAssetIds,
|
|
savedOriginalUrls,
|
|
};
|
|
}
|
|
|
|
const modelConfig = await resolveBillingModelByReference(modelId);
|
|
const totalCost = savedImageSizes.reduce((sum, imageSize) => {
|
|
return sum + getModelCreditCost(modelConfig, "image", { imageSize });
|
|
}, 0);
|
|
if (totalCost <= 0) {
|
|
return {
|
|
savedImageCount,
|
|
savedImageIds,
|
|
savedAssetIds,
|
|
savedOriginalUrls,
|
|
};
|
|
}
|
|
|
|
const consumeResult = await consumeUserPoints({
|
|
userId,
|
|
amount: totalCost,
|
|
reason: "image_generation",
|
|
modelId: modelConfig?.id || modelId,
|
|
chatId: chatId ?? undefined,
|
|
messageId: messageId ?? undefined,
|
|
note: `模型(${modelConfig?.displayName || modelId}) 生图 ${savedImageCount} 张`,
|
|
});
|
|
|
|
if (!consumeResult.ok) {
|
|
console.warn(
|
|
`[Points] Failed to consume points after generation. user=${userId} model=${modelId} cost=${totalCost}`,
|
|
);
|
|
} else if (consumeResult.skipped) {
|
|
console.log(
|
|
`[Points] Skipped ${totalCost} image points for VIP user=${userId} model=${modelId}`,
|
|
);
|
|
} else {
|
|
console.log(
|
|
`[Points] Consumed ${totalCost} points for ${savedImageCount} image(s). balance=${consumeResult.balance}`,
|
|
);
|
|
}
|
|
} catch (error) {
|
|
console.error("[Points] Failed to consume points after image generation:", error);
|
|
}
|
|
|
|
return {
|
|
savedImageCount,
|
|
savedImageIds,
|
|
savedAssetIds,
|
|
savedOriginalUrls,
|
|
};
|
|
}
|