ai-chatbot/lib/ai/providers/ruxa-image.ts

600 lines
18 KiB
TypeScript
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

import "server-only";
import sharp from "sharp";
import { createDeferredUploadedImage } from "@/lib/storage/image-archive";
import { uploadFile } from "@/lib/storage/service";
import {
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({
modelConfig,
}: {
modelConfig?: RuxaImageModelConfig | null;
}) {
if (modelConfig?.ruxaTaskInputMode === "size-image") {
return "size-image";
}
if (modelConfig?.ruxaTaskInputMode === "resolution-based") {
return "resolution-based";
}
return "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 deferred = await createDeferredUploadedImage(imageUrl, {
providerType: "ruxa-image",
});
if (deferred) return deferred;
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) => `![thumb|||${image.original}|||${resolution}](${image.thumbnail})`)
.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({
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,
taskInputMode,
imageInputCount: imageUrls.length,
inputKeys: Object.keys(input),
});
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,
};
},
}),
};
}