598 lines
18 KiB
TypeScript
598 lines
18 KiB
TypeScript
import "server-only";
|
||
|
||
import sharp from "sharp";
|
||
import { uploadFile } from "@/lib/storage/service";
|
||
import {
|
||
RUXA_NANO_BANANA_2_RUNTIME_MODEL_ID,
|
||
resolveRuxaImageRuntimeModelId,
|
||
} from "@/lib/ai/ruxa-image-models";
|
||
|
||
interface RuxaImageConfig {
|
||
apiKey: string;
|
||
baseURL: string;
|
||
}
|
||
|
||
interface RuxaImageParams {
|
||
aspectRatio?: string;
|
||
imageSize?: string;
|
||
imageCount?: number;
|
||
outputFormat?: "png" | "jpg" | "jpeg" | "webp";
|
||
}
|
||
|
||
interface RuxaImageModelConfig {
|
||
supportsSize?: boolean;
|
||
supports1K?: boolean;
|
||
supports2K?: boolean;
|
||
supports4K?: boolean;
|
||
ruxaTaskInputMode?: "size-image" | "resolution-based";
|
||
}
|
||
|
||
interface UploadedImage {
|
||
thumbnail: string;
|
||
original: string;
|
||
}
|
||
|
||
const DEFAULT_RESOLUTION = "1K";
|
||
const DEFAULT_OUTPUT_FORMAT = "png";
|
||
const POLL_INTERVAL_MS = 2000;
|
||
const POLL_MAX_ATTEMPTS = 240;
|
||
|
||
function normalizeRuxaModelId(modelId?: string) {
|
||
const raw = String(modelId || "").trim();
|
||
if (!raw) {
|
||
throw new Error("RUXA 模型 ID 未配置");
|
||
}
|
||
return raw;
|
||
}
|
||
|
||
function resolveRuxaRuntimeModelId(modelId: string) {
|
||
const normalized = modelId.trim();
|
||
return resolveRuxaImageRuntimeModelId(normalized);
|
||
}
|
||
|
||
function normalizeUrlForCompare(url: string) {
|
||
try {
|
||
const parsed = new URL(url);
|
||
parsed.search = "";
|
||
parsed.hash = "";
|
||
return parsed.toString();
|
||
} catch {
|
||
return url.split("?")[0]?.split("#")[0] || url;
|
||
}
|
||
}
|
||
|
||
function filterSourceImages(urls: string[], sourceUrls: string[]) {
|
||
if (sourceUrls.length === 0) return urls;
|
||
const sourceSet = new Set(sourceUrls.map((url) => normalizeUrlForCompare(url)));
|
||
return urls.filter((url) => !sourceSet.has(normalizeUrlForCompare(url)));
|
||
}
|
||
|
||
function joinApiPath(baseURL: string, path: string) {
|
||
const cleanBase = baseURL.replace(/\/+$/, "");
|
||
const cleanPath = path.startsWith("/") ? path : `/${path}`;
|
||
return `${cleanBase}${cleanPath}`;
|
||
}
|
||
|
||
function normalizeAspectRatio(aspectRatio?: string) {
|
||
if (!aspectRatio) return undefined;
|
||
const normalized = aspectRatio.trim().toLowerCase();
|
||
if (normalized === "auto") return undefined;
|
||
const ratio = aspectRatio.trim();
|
||
if (
|
||
["1:1", "16:9", "9:16", "4:3", "3:4", "2:3", "3:2", "5:4", "4:5", "21:9"].includes(ratio)
|
||
) {
|
||
return ratio;
|
||
}
|
||
return undefined;
|
||
}
|
||
|
||
function normalizeResolution(imageSize?: string) {
|
||
if (!imageSize) return DEFAULT_RESOLUTION;
|
||
const normalized = imageSize.trim().toUpperCase();
|
||
if (["1K", "2K", "4K"].includes(normalized)) return normalized;
|
||
return DEFAULT_RESOLUTION;
|
||
}
|
||
|
||
function normalizeOutputFormat(outputFormat?: string): "png" | "jpg" {
|
||
if (!outputFormat) return DEFAULT_OUTPUT_FORMAT;
|
||
const normalized = outputFormat.trim().toLowerCase();
|
||
if (normalized === "jpg" || normalized === "jpeg") return "jpg";
|
||
return "png";
|
||
}
|
||
|
||
function isResolutionSupported(resolution: string, config?: RuxaImageModelConfig | null) {
|
||
if (config?.supportsSize === false) {
|
||
return false;
|
||
}
|
||
|
||
const normalized = resolution.toLowerCase();
|
||
if (normalized === "1k") return config?.supports1K !== false;
|
||
if (normalized === "2k") return config?.supports2K !== false;
|
||
if (normalized === "4k") return config?.supports4K !== false;
|
||
return true;
|
||
}
|
||
|
||
function resolveRuxaTaskInputMode({
|
||
activeModelId,
|
||
modelConfig,
|
||
}: {
|
||
activeModelId: string;
|
||
modelConfig?: RuxaImageModelConfig | null;
|
||
}) {
|
||
if (modelConfig?.ruxaTaskInputMode === "size-image") {
|
||
return "size-image";
|
||
}
|
||
|
||
if (modelConfig?.ruxaTaskInputMode === "resolution-based") {
|
||
return "resolution-based";
|
||
}
|
||
|
||
return activeModelId === RUXA_NANO_BANANA_2_RUNTIME_MODEL_ID
|
||
? "size-image"
|
||
: "resolution-based";
|
||
}
|
||
|
||
function buildRuxaTaskInput({
|
||
prompt,
|
||
imageUrls,
|
||
aspectRatio,
|
||
resolution,
|
||
outputFormat,
|
||
modelConfig,
|
||
taskInputMode,
|
||
}: {
|
||
prompt: string;
|
||
imageUrls: string[];
|
||
aspectRatio?: string;
|
||
resolution: string;
|
||
outputFormat: string;
|
||
modelConfig?: RuxaImageModelConfig | null;
|
||
taskInputMode: "size-image" | "resolution-based";
|
||
}) {
|
||
const input: Record<string, unknown> = {
|
||
prompt,
|
||
output_format: outputFormat,
|
||
};
|
||
const ratio = aspectRatio || "1:1";
|
||
|
||
input.size = ratio;
|
||
if (taskInputMode === "size-image") {
|
||
if (imageUrls[0]) {
|
||
input.image = imageUrls[0];
|
||
}
|
||
} else {
|
||
input.aspect_ratio = ratio;
|
||
if (isResolutionSupported(resolution, modelConfig)) {
|
||
input.resolution = resolution;
|
||
}
|
||
if (imageUrls.length > 0) {
|
||
input.image_input = imageUrls.slice(0, 8);
|
||
}
|
||
}
|
||
|
||
return input;
|
||
}
|
||
|
||
function extractTaskId(payload: any): string | undefined {
|
||
if (typeof payload === "string" && payload.trim()) {
|
||
return payload.trim();
|
||
}
|
||
if (typeof payload?.data === "string" && payload.data.trim()) {
|
||
return payload.data.trim();
|
||
}
|
||
return (
|
||
payload?.taskId ||
|
||
payload?.task_id ||
|
||
payload?.id ||
|
||
payload?.task?.id ||
|
||
payload?.data?.taskId ||
|
||
payload?.data?.task_id ||
|
||
payload?.data?.id ||
|
||
payload?.data?.task?.id
|
||
);
|
||
}
|
||
|
||
function extractStatus(payload: any): string | undefined {
|
||
const raw =
|
||
payload?.status ||
|
||
payload?.taskStatus ||
|
||
payload?.task_status ||
|
||
payload?.state ||
|
||
payload?.data?.status ||
|
||
payload?.data?.taskStatus ||
|
||
payload?.data?.task_status ||
|
||
payload?.data?.state;
|
||
if (typeof raw !== "string") return undefined;
|
||
return raw.toLowerCase();
|
||
}
|
||
|
||
function extractErrorMessage(payload: any): string {
|
||
let parsedResultJson: any = null;
|
||
const rawResultJson =
|
||
typeof payload?.resultJson === "string" && payload.resultJson.trim()
|
||
? payload.resultJson
|
||
: typeof payload?.data?.resultJson === "string" && payload.data.resultJson.trim()
|
||
? payload.data.resultJson
|
||
: "";
|
||
|
||
if (rawResultJson) {
|
||
try {
|
||
parsedResultJson = JSON.parse(rawResultJson);
|
||
} catch {
|
||
parsedResultJson = null;
|
||
}
|
||
}
|
||
|
||
return (
|
||
payload?.failMsg ||
|
||
payload?.data?.failMsg ||
|
||
payload?.message ||
|
||
payload?.error ||
|
||
payload?.errorMessage ||
|
||
parsedResultJson?.failMsg ||
|
||
parsedResultJson?.error ||
|
||
rawResultJson ||
|
||
payload?.data?.message ||
|
||
payload?.data?.error ||
|
||
"RUXA 任务执行失败"
|
||
);
|
||
}
|
||
|
||
function compactPayloadSnippet(payload: any): string {
|
||
try {
|
||
const json = JSON.stringify(payload);
|
||
return json.length > 320 ? `${json.slice(0, 320)}...` : json;
|
||
} catch {
|
||
return String(payload || "");
|
||
}
|
||
}
|
||
|
||
function extractPromptAndImages(messages: Array<{ role: string; content: any }>) {
|
||
const lastUserMessage = messages.filter((message) => message.role === "user").pop();
|
||
if (!lastUserMessage) {
|
||
throw new Error("No user message found");
|
||
}
|
||
|
||
let prompt = "";
|
||
const imageUrls: string[] = [];
|
||
|
||
const pushImageUrl = (url?: string) => {
|
||
if (!url) return;
|
||
const normalized = String(url).trim();
|
||
if (normalized.startsWith("data:")) {
|
||
throw new Error("RUXA 不支持 Base64 图片传输，请先上传图片后再编辑");
|
||
}
|
||
if (!normalized.startsWith("http")) return;
|
||
if (!imageUrls.includes(normalized)) {
|
||
imageUrls.push(normalized);
|
||
}
|
||
};
|
||
|
||
if (typeof lastUserMessage.content === "string") {
|
||
prompt = lastUserMessage.content.trim();
|
||
} else if (Array.isArray(lastUserMessage.content)) {
|
||
for (const part of lastUserMessage.content) {
|
||
if (part?.type === "text" && typeof part.text === "string") {
|
||
prompt = `${prompt} ${part.text}`.trim();
|
||
} else if (part?.type === "image_url" && part.image_url?.url) {
|
||
pushImageUrl(part.image_url.url);
|
||
} else if (part?.type === "file") {
|
||
pushImageUrl(part.url || part.data);
|
||
} else if (part?.type === "image") {
|
||
pushImageUrl(part.image || part.url);
|
||
}
|
||
}
|
||
}
|
||
|
||
if (!prompt) {
|
||
throw new Error("No prompt found in user message");
|
||
}
|
||
|
||
return { prompt, imageUrls };
|
||
}
|
||
|
||
function collectImageUrls(payload: any): string[] {
|
||
const urls = new Set<string>();
|
||
|
||
const isImageLikeKey = (key: string) =>
|
||
/image|img|url|output|result|artifact|file/i.test(key);
|
||
|
||
const isLikelyImageUrl = (value: string) =>
|
||
/^https?:\/\//i.test(value) &&
|
||
(/\.(png|jpg|jpeg|webp|gif)(\?|$)/i.test(value) || value.includes("image"));
|
||
|
||
const walk = (value: any, parentKey = "") => {
|
||
if (typeof value === "string") {
|
||
if (isLikelyImageUrl(value) || (isImageLikeKey(parentKey) && /^https?:\/\//i.test(value))) {
|
||
urls.add(value);
|
||
return;
|
||
}
|
||
|
||
const trimmed = value.trim();
|
||
if (
|
||
trimmed.length > 1 &&
|
||
trimmed.length < 200000 &&
|
||
(trimmed.startsWith("{") || trimmed.startsWith("["))
|
||
) {
|
||
try {
|
||
const parsed = JSON.parse(trimmed);
|
||
walk(parsed, parentKey);
|
||
} catch {
|
||
// ignore non-JSON strings
|
||
}
|
||
}
|
||
return;
|
||
}
|
||
|
||
if (Array.isArray(value)) {
|
||
for (const item of value) {
|
||
walk(item, parentKey);
|
||
}
|
||
return;
|
||
}
|
||
|
||
if (value && typeof value === "object") {
|
||
for (const [key, nestedValue] of Object.entries(value)) {
|
||
walk(nestedValue, key);
|
||
}
|
||
}
|
||
};
|
||
|
||
walk(payload);
|
||
return Array.from(urls);
|
||
}
|
||
|
||
function detectImageExtension(imageUrl: string, contentType?: string | null) {
|
||
const type = (contentType || "").toLowerCase();
|
||
if (type.includes("jpeg") || type.includes("jpg")) return "jpg";
|
||
if (type.includes("webp")) return "webp";
|
||
if (type.includes("png")) return "png";
|
||
|
||
const cleanUrl = imageUrl.split("?")[0]?.toLowerCase() || "";
|
||
if (cleanUrl.endsWith(".jpg") || cleanUrl.endsWith(".jpeg")) return "jpg";
|
||
if (cleanUrl.endsWith(".webp")) return "webp";
|
||
if (cleanUrl.endsWith(".png")) return "png";
|
||
return "png";
|
||
}
|
||
|
||
async function createThumbnail(buffer: Buffer): Promise<Buffer> {
|
||
try {
|
||
return await sharp(buffer)
|
||
.resize(400, 400, { fit: "inside", withoutEnlargement: true })
|
||
.jpeg({ quality: 80 })
|
||
.toBuffer();
|
||
} catch (error) {
|
||
console.error("[Ruxa Image] Failed to create thumbnail:", error);
|
||
return buffer;
|
||
}
|
||
}
|
||
|
||
function generateFilename(prefix: string, extension: string) {
|
||
const ts = Date.now();
|
||
const rand = Math.random().toString(36).slice(2, 8);
|
||
return `${prefix}-${ts}-${rand}.${extension}`;
|
||
}
|
||
|
||
async function downloadAndUploadImage(imageUrl: string): Promise<UploadedImage> {
|
||
const response = await fetch(imageUrl);
|
||
if (!response.ok) {
|
||
throw new Error(`Failed to download image: ${response.status}`);
|
||
}
|
||
|
||
const arrayBuffer = await response.arrayBuffer();
|
||
const buffer = Buffer.from(arrayBuffer);
|
||
const thumbnailBuffer = await createThumbnail(buffer);
|
||
|
||
const originalExt = detectImageExtension(imageUrl, response.headers.get("content-type"));
|
||
const originalFilename = generateFilename("ruxa-image-orig", originalExt);
|
||
const originalArrayBuffer = buffer.buffer.slice(
|
||
buffer.byteOffset,
|
||
buffer.byteOffset + buffer.byteLength
|
||
) as ArrayBuffer;
|
||
const originalUploaded = await uploadFile(originalFilename, originalArrayBuffer);
|
||
|
||
const thumbFilename = generateFilename("ruxa-image-thumb", "jpg");
|
||
const thumbArrayBuffer = thumbnailBuffer.buffer.slice(
|
||
thumbnailBuffer.byteOffset,
|
||
thumbnailBuffer.byteOffset + thumbnailBuffer.byteLength
|
||
) as ArrayBuffer;
|
||
const thumbUploaded = await uploadFile(thumbFilename, thumbArrayBuffer);
|
||
|
||
return {
|
||
thumbnail: thumbUploaded.url,
|
||
original: originalUploaded.url,
|
||
};
|
||
}
|
||
|
||
async function waitForTaskResult(baseURL: string, apiKey: string, taskId: string, sourceUrls: string[]) {
|
||
const endpointCandidates = [
|
||
(id: string) => joinApiPath(baseURL, `/api/v1/tasks/query/${id}`),
|
||
(id: string) => joinApiPath(baseURL, `/api/v1/tasks/${id}`),
|
||
(id: string) => `${joinApiPath(baseURL, "/api/v1/tasks/query")}?taskId=${encodeURIComponent(id)}`,
|
||
];
|
||
let endpointIndex = 0;
|
||
|
||
for (let attempt = 1; attempt <= POLL_MAX_ATTEMPTS; attempt++) {
|
||
try {
|
||
const endpoint = endpointCandidates[endpointIndex]!(taskId);
|
||
const response = await fetch(endpoint, {
|
||
headers: {
|
||
Authorization: `Bearer ${apiKey}`,
|
||
"Content-Type": "application/json",
|
||
},
|
||
});
|
||
|
||
if (!response.ok) {
|
||
if (response.status === 404 && endpointIndex < endpointCandidates.length - 1) {
|
||
endpointIndex++;
|
||
}
|
||
await new Promise((resolve) => setTimeout(resolve, POLL_INTERVAL_MS));
|
||
continue;
|
||
}
|
||
|
||
const payload = await response.json();
|
||
const imageUrls = filterSourceImages(collectImageUrls(payload), sourceUrls);
|
||
const status = extractStatus(payload);
|
||
|
||
if (imageUrls.length > 0) {
|
||
return imageUrls;
|
||
}
|
||
|
||
if (status && ["failed", "fail", "error", "cancelled", "canceled"].includes(status)) {
|
||
throw new Error(extractErrorMessage(payload));
|
||
}
|
||
|
||
if (status && ["success", "succeeded", "completed", "done", "finished"].includes(status)) {
|
||
throw new Error("任务已完成但未返回图片地址");
|
||
}
|
||
} catch (error) {
|
||
if (error instanceof Error && error.message !== "任务已完成但未返回图片地址") {
|
||
throw error;
|
||
}
|
||
}
|
||
|
||
await new Promise((resolve) => setTimeout(resolve, POLL_INTERVAL_MS));
|
||
}
|
||
|
||
throw new Error("RUXA 任务超时，请稍后重试");
|
||
}
|
||
|
||
function formatMarkdown(images: UploadedImage[], resolution: string) {
|
||
return images
|
||
.map((image) => ``)
|
||
.join("\n\n");
|
||
}
|
||
|
||
export function createRuxaImage({ apiKey, baseURL }: RuxaImageConfig) {
|
||
return {
|
||
provider: "ruxa-image",
|
||
chat: (modelId: string) => ({
|
||
provider: "ruxa-image",
|
||
doGenerate: async (options: {
|
||
messages: Array<{ role: string; content: any }>;
|
||
imageParams?: RuxaImageParams;
|
||
imageGenConfig?: RuxaImageModelConfig | null;
|
||
overrideModelId?: string;
|
||
}) => {
|
||
const { prompt, imageUrls } = extractPromptAndImages(options.messages);
|
||
const aspectRatio = normalizeAspectRatio(options.imageParams?.aspectRatio);
|
||
const resolution = normalizeResolution(options.imageParams?.imageSize);
|
||
const requestedModelId = options.overrideModelId || modelId;
|
||
const normalizedRequestedModelId = normalizeRuxaModelId(requestedModelId);
|
||
const activeModelId = resolveRuxaRuntimeModelId(normalizedRequestedModelId);
|
||
if (normalizedRequestedModelId !== activeModelId) {
|
||
console.warn("[Ruxa Image] 运行时模型已映射", {
|
||
requestedModelId: normalizedRequestedModelId,
|
||
runtimeModelId: activeModelId,
|
||
});
|
||
console.warn("[Ruxa Image] 请确认数据库中已同步更新该模型的运行时 endpoint", {
|
||
requestedModelId: normalizedRequestedModelId,
|
||
runtimeModelId: activeModelId,
|
||
});
|
||
}
|
||
const taskInputMode = resolveRuxaTaskInputMode({
|
||
activeModelId,
|
||
modelConfig: options.imageGenConfig,
|
||
});
|
||
const outputFormat = normalizeOutputFormat(options.imageParams?.outputFormat);
|
||
|
||
const createEndpoint = joinApiPath(baseURL, "/api/v1/tasks/create");
|
||
const callbackUrl = process.env.RUXA_CALLBACK_URL?.trim();
|
||
const input = buildRuxaTaskInput({
|
||
prompt,
|
||
imageUrls,
|
||
aspectRatio,
|
||
resolution,
|
||
outputFormat,
|
||
modelConfig: options.imageGenConfig,
|
||
taskInputMode,
|
||
});
|
||
const requestBody = {
|
||
input,
|
||
model: activeModelId,
|
||
...(callbackUrl ? { callback_url: callbackUrl } : {}),
|
||
};
|
||
console.log("[Ruxa Image] Creating task:", {
|
||
model: activeModelId,
|
||
resolution: input.resolution || null,
|
||
aspectRatio: input.aspect_ratio || input.size || null,
|
||
outputFormat,
|
||
imageInputCount: imageUrls.length,
|
||
});
|
||
const createResponse = await fetch(createEndpoint, {
|
||
method: "POST",
|
||
headers: {
|
||
Authorization: `Bearer ${apiKey}`,
|
||
"Content-Type": "application/json",
|
||
},
|
||
body: JSON.stringify(requestBody),
|
||
});
|
||
const createRawText = await createResponse.text();
|
||
let createPayload: any = null;
|
||
try {
|
||
createPayload = JSON.parse(createRawText);
|
||
} catch {
|
||
createPayload = createRawText;
|
||
}
|
||
if (!createResponse.ok) {
|
||
throw new Error(
|
||
`${extractErrorMessage(createPayload)} (HTTP ${createResponse.status})`
|
||
);
|
||
}
|
||
|
||
const createCode =
|
||
typeof createPayload?.code === "number"
|
||
? createPayload.code
|
||
: Number.isFinite(Number(createPayload?.code))
|
||
? Number(createPayload.code)
|
||
: null;
|
||
if (createCode !== null && createCode !== 200) {
|
||
throw new Error(extractErrorMessage(createPayload));
|
||
}
|
||
|
||
let outputUrls = filterSourceImages(collectImageUrls(createPayload), imageUrls);
|
||
if (outputUrls.length === 0) {
|
||
const taskId = extractTaskId(createPayload);
|
||
if (!taskId) {
|
||
throw new Error(
|
||
`RUXA 创建任务成功，但未返回 taskId，响应片段: ${compactPayloadSnippet(createPayload)}`
|
||
);
|
||
}
|
||
outputUrls = await waitForTaskResult(baseURL, apiKey, taskId, imageUrls);
|
||
}
|
||
|
||
const uploadedImages: UploadedImage[] = [];
|
||
for (const outputUrl of outputUrls.slice(0, options.imageParams?.imageCount || 1)) {
|
||
uploadedImages.push(await downloadAndUploadImage(outputUrl));
|
||
}
|
||
|
||
if (uploadedImages.length === 0) {
|
||
throw new Error("RUXA 未返回可用图片");
|
||
}
|
||
|
||
const text = formatMarkdown(uploadedImages, resolution);
|
||
return {
|
||
text,
|
||
finishReason: "stop" as const,
|
||
usage: {
|
||
promptTokens: 0,
|
||
completionTokens: text.length,
|
||
},
|
||
};
|
||
},
|
||
doStream: async function* (options: {
|
||
messages: Array<{ role: string; content: any }>;
|
||
imageParams?: RuxaImageParams;
|
||
imageGenConfig?: RuxaImageModelConfig | null;
|
||
overrideModelId?: string;
|
||
}) {
|
||
const result = await this.doGenerate(options);
|
||
yield {
|
||
type: "text-delta",
|
||
textDelta: result.text,
|
||
};
|
||
},
|
||
}),
|
||
};
|
||
}
|