455 lines
15 KiB
TypeScript
455 lines
15 KiB
TypeScript
import "server-only";
|
||
|
||
import sharp from "sharp";
|
||
import { ATLASCLOUD_BASE_URL } from "@/lib/ai/atlascloud-models";
|
||
import {
|
||
createAtlasCloudRequestBody,
|
||
type AtlasCloudImageParams,
|
||
type AtlasCloudModelConfig,
|
||
} from "@/lib/ai/atlascloud-request";
|
||
import { createDeferredUploadedImage } from "@/lib/storage/image-archive";
|
||
import { uploadFile } from "@/lib/storage/service";
|
||
|
||
interface AtlasCloudConfig {
|
||
apiKey: string;
|
||
baseURL?: string;
|
||
}
|
||
|
||
interface UploadedImage {
|
||
thumbnail: string;
|
||
original: string;
|
||
}
|
||
|
||
const POLL_INTERVAL_MS = 3000;
|
||
const POLL_TIMEOUT_MS = 20 * 60 * 1000;
|
||
const DOWNLOAD_TIMEOUT_MS = 120_000;
|
||
const TRANSIENT_FAILURE_LIMIT = 8;
|
||
|
||
function normalizeBaseUrl(baseURL?: string) {
|
||
return String(baseURL || ATLASCLOUD_BASE_URL).replace(/\/+$/, "");
|
||
}
|
||
|
||
function joinApiPath(baseURL: string | undefined, path: string) {
|
||
const base = normalizeBaseUrl(baseURL);
|
||
const cleanPath = path.startsWith("/") ? path : `/${path}`;
|
||
if (base.endsWith("/api/v1") && cleanPath.startsWith("/api/v1/")) {
|
||
return `${base}${cleanPath.slice("/api/v1".length)}`;
|
||
}
|
||
return `${base}${cleanPath}`;
|
||
}
|
||
|
||
function delay(ms: number) {
|
||
return new Promise((resolve) => setTimeout(resolve, ms));
|
||
}
|
||
|
||
function parseJsonResponse(response: Response) {
|
||
return response.text().then((text) => {
|
||
if (!text.trim()) return null;
|
||
try {
|
||
return JSON.parse(text);
|
||
} catch {
|
||
return text;
|
||
}
|
||
});
|
||
}
|
||
|
||
function extractErrorMessage(payload: any, fallback = "Atlas Cloud 请求失败") {
|
||
if (typeof payload === "string" && payload.trim()) return payload.trim();
|
||
const value =
|
||
payload?.data?.error?.message ||
|
||
payload?.data?.error ||
|
||
payload?.error?.message ||
|
||
payload?.error ||
|
||
payload?.message ||
|
||
fallback;
|
||
return typeof value === "string" && value.trim() ? value.trim() : fallback;
|
||
}
|
||
|
||
function unwrapPrediction(payload: any) {
|
||
return payload?.data && typeof payload.data === "object" ? payload.data : payload;
|
||
}
|
||
|
||
function extractTaskId(payload: any) {
|
||
const prediction = unwrapPrediction(payload);
|
||
return String(
|
||
prediction?.id || prediction?.request_id || prediction?.task_id || ""
|
||
).trim();
|
||
}
|
||
|
||
function extractOutputUrls(payload: any) {
|
||
const prediction = unwrapPrediction(payload);
|
||
const candidates = [
|
||
prediction?.outputs,
|
||
prediction?.output,
|
||
prediction?.images,
|
||
prediction?.urls?.outputs,
|
||
];
|
||
const urls: string[] = [];
|
||
for (const candidate of candidates) {
|
||
const values = Array.isArray(candidate) ? candidate : candidate ? [candidate] : [];
|
||
for (const value of values) {
|
||
const url =
|
||
typeof value === "string"
|
||
? value
|
||
: String(value?.url || value?.image_url || value?.src || "");
|
||
if (/^(?:https?:\/\/|data:image\/)/i.test(url) && !urls.includes(url)) {
|
||
urls.push(url);
|
||
}
|
||
}
|
||
}
|
||
return urls;
|
||
}
|
||
|
||
function extractPromptAndMedia(
|
||
messages: Array<{ role: string; content?: any; parts?: any[] }>
|
||
) {
|
||
const message = messages.filter((item) => item.role === "user").pop();
|
||
if (!message) throw new Error("未找到用户消息");
|
||
const texts: string[] = [];
|
||
const imageUrls: string[] = [];
|
||
const videoUrls: string[] = [];
|
||
|
||
const pushText = (value: unknown) => {
|
||
if (typeof value === "string" && value.trim()) texts.push(value.trim());
|
||
};
|
||
const pushMedia = (value: unknown, mediaType?: unknown) => {
|
||
const url = String(
|
||
typeof value === "string"
|
||
? value
|
||
: (value as any)?.url ||
|
||
(value as any)?.imageUrl ||
|
||
(value as any)?.fileUrl ||
|
||
(value as any)?.src ||
|
||
""
|
||
).trim();
|
||
if (!url || !/^(?:https?:\/\/|data:)/i.test(url)) return;
|
||
const normalizedMediaType = String(mediaType || "").toLowerCase();
|
||
const isVideo =
|
||
normalizedMediaType.startsWith("video/") ||
|
||
url.startsWith("data:video/") ||
|
||
/\.(?:mp4|mov|webm|m4v)(?:[?#].*)?$/i.test(url);
|
||
const target = isVideo ? videoUrls : imageUrls;
|
||
if (!target.includes(url)) target.push(url);
|
||
};
|
||
const visitPart = (part: any) => {
|
||
if (!part) return;
|
||
if (typeof part === "string") return pushText(part);
|
||
if (part.type === "text" || typeof part.text === "string") pushText(part.text);
|
||
if (part.type === "image_url") pushMedia(part.image_url, "image/jpeg");
|
||
if (part.type === "image") pushMedia(part.image || part.url, "image/jpeg");
|
||
if (part.type === "file") {
|
||
pushMedia(
|
||
part.url || part.file || part.data,
|
||
part.mediaType || part.contentType || part.mimeType
|
||
);
|
||
}
|
||
if (part.inlineData?.data) {
|
||
pushMedia(
|
||
`data:${part.inlineData.mimeType || "image/png"};base64,${part.inlineData.data}`,
|
||
part.inlineData.mimeType
|
||
);
|
||
}
|
||
};
|
||
|
||
if (typeof message.content === "string") pushText(message.content);
|
||
else if (Array.isArray(message.content)) message.content.forEach(visitPart);
|
||
if (Array.isArray(message.parts)) message.parts.forEach(visitPart);
|
||
|
||
const prompt = texts.join("\n").trim();
|
||
if (!prompt) throw new Error("请输入图片生成或编辑要求");
|
||
return { prompt, imageUrls, videoUrls };
|
||
}
|
||
|
||
function parseDataUrl(value: string) {
|
||
const match = value.match(/^data:([^;,]+);base64,(.+)$/s);
|
||
if (!match) throw new Error("Atlas Cloud 参考素材 Base64 格式无效");
|
||
return {
|
||
mimeType: match[1] || "image/png",
|
||
buffer: Buffer.from(match[2] || "", "base64"),
|
||
};
|
||
}
|
||
|
||
async function loadMediaSource(mediaUrl: string) {
|
||
if (mediaUrl.startsWith("data:")) return parseDataUrl(mediaUrl);
|
||
const response = await fetch(mediaUrl, {
|
||
cache: "no-store",
|
||
redirect: "follow",
|
||
signal: AbortSignal.timeout(DOWNLOAD_TIMEOUT_MS),
|
||
headers: {
|
||
Accept: "image/*,video/*,*/*;q=0.8",
|
||
"User-Agent": "yedir-ai-atlascloud/1.0",
|
||
},
|
||
});
|
||
if (!response.ok) {
|
||
throw new Error(`Atlas Cloud 参考素材下载失败（HTTP ${response.status}）`);
|
||
}
|
||
return {
|
||
mimeType: response.headers.get("content-type")?.split(";")[0] || "image/png",
|
||
buffer: Buffer.from(await response.arrayBuffer()),
|
||
};
|
||
}
|
||
|
||
function detectExtension(mimeType: string) {
|
||
if (/jpe?g/i.test(mimeType)) return "jpg";
|
||
if (/webp/i.test(mimeType)) return "webp";
|
||
if (/gif/i.test(mimeType)) return "gif";
|
||
if (/quicktime/i.test(mimeType)) return "mov";
|
||
if (/webm/i.test(mimeType)) return "webm";
|
||
if (/mp4/i.test(mimeType)) return "mp4";
|
||
return "png";
|
||
}
|
||
|
||
async function uploadReferenceMedia(
|
||
baseURL: string | undefined,
|
||
apiKey: string,
|
||
mediaUrl: string
|
||
) {
|
||
if (/^https:\/\//i.test(mediaUrl)) return mediaUrl;
|
||
const source = await loadMediaSource(mediaUrl);
|
||
const form = new FormData();
|
||
form.append(
|
||
"file",
|
||
new Blob([source.buffer], { type: source.mimeType }),
|
||
`reference-${Date.now()}.${detectExtension(source.mimeType)}`
|
||
);
|
||
const response = await fetch(joinApiPath(baseURL, "/api/v1/model/uploadMedia"), {
|
||
method: "POST",
|
||
headers: { Authorization: `Bearer ${apiKey}` },
|
||
body: form,
|
||
signal: AbortSignal.timeout(DOWNLOAD_TIMEOUT_MS),
|
||
});
|
||
const payload = await parseJsonResponse(response);
|
||
if (!response.ok) {
|
||
throw new Error(
|
||
`${extractErrorMessage(payload, "Atlas Cloud 参考素材上传失败")}（HTTP ${response.status}）`
|
||
);
|
||
}
|
||
const data = payload?.data && typeof payload.data === "object" ? payload.data : payload;
|
||
const url = String(
|
||
data?.url || data?.file_url || data?.fileUrl || data?.media_url || data?.mediaUrl || ""
|
||
).trim();
|
||
if (!url) throw new Error("Atlas Cloud 参考素材上传成功但未返回文件地址");
|
||
return url;
|
||
}
|
||
|
||
async function submitTask(params: {
|
||
apiKey: string;
|
||
baseURL?: string;
|
||
endpoint: string;
|
||
body: Record<string, unknown>;
|
||
}) {
|
||
const response = await fetch(joinApiPath(params.baseURL, params.endpoint), {
|
||
method: "POST",
|
||
headers: {
|
||
Authorization: `Bearer ${params.apiKey}`,
|
||
"Content-Type": "application/json",
|
||
},
|
||
body: JSON.stringify(params.body),
|
||
signal: AbortSignal.timeout(120_000),
|
||
});
|
||
const payload = await parseJsonResponse(response);
|
||
if (!response.ok) {
|
||
throw new Error(
|
||
`${extractErrorMessage(payload)}（HTTP ${response.status}）`
|
||
);
|
||
}
|
||
const immediateOutputs = extractOutputUrls(payload);
|
||
if (immediateOutputs.length > 0) return { outputs: immediateOutputs, taskId: "" };
|
||
const taskId = extractTaskId(payload);
|
||
if (!taskId) throw new Error("Atlas Cloud 未返回任务 ID");
|
||
return { outputs: [] as string[], taskId };
|
||
}
|
||
|
||
async function waitForTask(params: {
|
||
apiKey: string;
|
||
baseURL?: string;
|
||
endpoint: string;
|
||
taskId: string;
|
||
}) {
|
||
const startedAt = Date.now();
|
||
let transientFailures = 0;
|
||
while (Date.now() - startedAt < POLL_TIMEOUT_MS) {
|
||
await delay(POLL_INTERVAL_MS);
|
||
try {
|
||
const endpoint = params.endpoint.replace(
|
||
/\{(?:request_id|id|task_id)\}/g,
|
||
encodeURIComponent(params.taskId)
|
||
);
|
||
const response = await fetch(joinApiPath(params.baseURL, endpoint), {
|
||
cache: "no-store",
|
||
headers: { Authorization: `Bearer ${params.apiKey}` },
|
||
signal: AbortSignal.timeout(60_000),
|
||
});
|
||
const payload = await parseJsonResponse(response);
|
||
if (!response.ok) {
|
||
if (response.status >= 500 || response.status === 429) {
|
||
throw new Error(`HTTP ${response.status}`);
|
||
}
|
||
throw new Error(extractErrorMessage(payload));
|
||
}
|
||
transientFailures = 0;
|
||
const prediction = unwrapPrediction(payload);
|
||
const status = String(prediction?.status || "").toLowerCase();
|
||
const outputs = extractOutputUrls(payload);
|
||
if (status === "completed" || status === "succeeded" || status === "success") {
|
||
if (outputs.length === 0) {
|
||
throw new Error("Atlas Cloud 任务完成但未返回图片");
|
||
}
|
||
return outputs;
|
||
}
|
||
if (["failed", "error", "timeout", "canceled", "cancelled"].includes(status)) {
|
||
throw new Error(extractErrorMessage(payload, "Atlas Cloud 图片生成失败"));
|
||
}
|
||
} catch (error) {
|
||
const message = error instanceof Error ? error.message : String(error);
|
||
if (!/^HTTP (?:429|5\d\d)$/.test(message) && !/fetch|network|abort/i.test(message)) {
|
||
throw error;
|
||
}
|
||
transientFailures += 1;
|
||
if (transientFailures >= TRANSIENT_FAILURE_LIMIT) throw error;
|
||
}
|
||
}
|
||
throw new Error("Atlas Cloud 图片生成超时，任务可能仍在处理，请稍后重试");
|
||
}
|
||
|
||
function toArrayBuffer(buffer: Buffer) {
|
||
return buffer.buffer.slice(
|
||
buffer.byteOffset,
|
||
buffer.byteOffset + buffer.byteLength
|
||
) as ArrayBuffer;
|
||
}
|
||
|
||
async function uploadImageBuffer(buffer: Buffer, mimeType: string) {
|
||
const extension = detectExtension(mimeType);
|
||
const thumbnailBuffer = await sharp(buffer, { failOn: "none" })
|
||
.resize(400, 400, { fit: "inside", withoutEnlargement: true })
|
||
.jpeg({ quality: 82 })
|
||
.toBuffer()
|
||
.catch(() => buffer);
|
||
const suffix = `${Date.now()}-${Math.random().toString(36).slice(2, 8)}`;
|
||
const [original, thumbnail] = await Promise.all([
|
||
uploadFile(`atlascloud-orig-${suffix}.${extension}`, toArrayBuffer(buffer)),
|
||
uploadFile(`atlascloud-thumb-${suffix}.jpg`, toArrayBuffer(thumbnailBuffer)),
|
||
]);
|
||
return { original: original.url, thumbnail: thumbnail.url };
|
||
}
|
||
|
||
async function archiveOutputImage(imageUrl: string): Promise<UploadedImage> {
|
||
if (imageUrl.startsWith("data:")) {
|
||
const parsed = parseDataUrl(imageUrl);
|
||
return uploadImageBuffer(parsed.buffer, parsed.mimeType);
|
||
}
|
||
const deferred = await createDeferredUploadedImage(imageUrl, {
|
||
providerType: "atlascloud",
|
||
});
|
||
if (deferred) return deferred;
|
||
const response = await fetch(imageUrl, {
|
||
cache: "no-store",
|
||
redirect: "follow",
|
||
signal: AbortSignal.timeout(DOWNLOAD_TIMEOUT_MS),
|
||
headers: { "User-Agent": "yedir-ai-atlascloud/1.0" },
|
||
});
|
||
if (!response.ok) {
|
||
throw new Error(`Atlas Cloud 结果图片下载失败（HTTP ${response.status}）`);
|
||
}
|
||
return uploadImageBuffer(
|
||
Buffer.from(await response.arrayBuffer()),
|
||
response.headers.get("content-type") || "image/png"
|
||
);
|
||
}
|
||
|
||
function formatMarkdown(images: UploadedImage[], imageSize: string) {
|
||
return images
|
||
.map((image) => ``)
|
||
.join("\n\n");
|
||
}
|
||
|
||
export function createAtlasCloud({ apiKey, baseURL }: AtlasCloudConfig) {
|
||
return {
|
||
provider: "atlascloud",
|
||
chat: (modelId: string, defaultModelConfig?: AtlasCloudModelConfig | null) => ({
|
||
provider: "atlascloud",
|
||
doGenerate: async (options: {
|
||
messages: Array<{ role: string; content?: any; parts?: any[] }>;
|
||
imageParams?: AtlasCloudImageParams;
|
||
imageGenConfig?: AtlasCloudModelConfig | null;
|
||
overrideModelId?: string;
|
||
}) => {
|
||
const runtimeModelId = String(options.overrideModelId || modelId || "").trim();
|
||
if (!runtimeModelId) throw new Error("Atlas Cloud 模型 ID 未配置");
|
||
const modelConfig = options.imageGenConfig || defaultModelConfig || {};
|
||
if (!modelConfig.atlasCloud) {
|
||
throw new Error("Atlas Cloud 模型能力配置缺失，请在系统设置中重新同步模型");
|
||
}
|
||
const { prompt, imageUrls, videoUrls } = extractPromptAndMedia(
|
||
options.messages
|
||
);
|
||
const maxReferenceImages = Math.max(
|
||
0,
|
||
Math.min(14, modelConfig.maxReferenceImages ?? modelConfig.atlasCloud.maxReferenceImages)
|
||
);
|
||
const referenceImages = await Promise.all(
|
||
imageUrls
|
||
.slice(0, maxReferenceImages)
|
||
.map((url) => uploadReferenceMedia(baseURL, apiKey, url))
|
||
);
|
||
const maxVideoReferences = Math.max(
|
||
0,
|
||
Math.min(
|
||
4,
|
||
modelConfig.maxVideoReferences ??
|
||
modelConfig.atlasCloud.maxVideoReferences ??
|
||
0
|
||
)
|
||
);
|
||
const referenceVideos = await Promise.all(
|
||
videoUrls
|
||
.slice(0, maxVideoReferences)
|
||
.map((url) => uploadReferenceMedia(baseURL, apiKey, url))
|
||
);
|
||
const { body, imageSize } = createAtlasCloudRequestBody({
|
||
modelId: runtimeModelId,
|
||
prompt,
|
||
referenceImages,
|
||
referenceVideos,
|
||
imageParams: options.imageParams,
|
||
modelConfig,
|
||
});
|
||
const submitted = await submitTask({
|
||
apiKey,
|
||
baseURL,
|
||
endpoint: modelConfig.atlasCloud.submitEndpoint,
|
||
body,
|
||
});
|
||
const outputUrls = submitted.outputs.length
|
||
? submitted.outputs
|
||
: await waitForTask({
|
||
apiKey,
|
||
baseURL,
|
||
endpoint: modelConfig.atlasCloud.pollEndpoint,
|
||
taskId: submitted.taskId,
|
||
});
|
||
const requestedCount = Math.max(1, options.imageParams?.imageCount || 1);
|
||
const archived = await Promise.all(
|
||
outputUrls.slice(0, requestedCount).map(archiveOutputImage)
|
||
);
|
||
const text = formatMarkdown(archived, imageSize);
|
||
return {
|
||
text,
|
||
finishReason: "stop" as const,
|
||
usage: { promptTokens: 0, completionTokens: text.length },
|
||
};
|
||
},
|
||
doStream: async function* (options: {
|
||
messages: Array<{ role: string; content?: any; parts?: any[] }>;
|
||
imageParams?: AtlasCloudImageParams;
|
||
imageGenConfig?: AtlasCloudModelConfig | null;
|
||
overrideModelId?: string;
|
||
}) {
|
||
const result = await this.doGenerate(options);
|
||
yield { type: "text-delta", textDelta: result.text };
|
||
},
|
||
}),
|
||
};
|
||
}
|