ai-chatbot/lib/ai/providers/midjourney.ts

466 lines
14 KiB
TypeScript

import "server-only";
import sharp from "sharp";
import { createDeferredUploadedImage } from "@/lib/storage/image-archive";
import { uploadFile } from "@/lib/storage/service";
interface MidjourneyConfig {
apiKey: string;
baseURL: string;
}
export interface MidjourneyImageParams {
aspectRatio?: string;
imageSize?: string;
imageCount?: number;
}
interface UploadedImage {
thumbnail: string;
original: string;
}
interface MidjourneyTaskButton {
customId?: string;
label?: string;
emoji?: string;
style?: number;
type?: number;
}
interface MidjourneyTaskResult {
id?: string;
imageUrl?: string;
status?: string;
progress?: string;
failReason?: string;
description?: string;
prompt?: string;
buttons?: MidjourneyTaskButton[];
}
const DEFAULT_TIMEOUT_MS = 600000;
const POLL_INTERVAL_MS = 3500;
const DEFAULT_IMAGE_SIZE = "2K";
function normalizeBaseURL(baseURL: string) {
return String(baseURL || "").trim().replace(/\/+$/, "");
}
function resolveMode(baseURL: string, modelId: string) {
const normalizedBaseURL = normalizeBaseURL(baseURL).toLowerCase();
const normalizedModelId = String(modelId || "").trim().toLowerCase();
if (
normalizedBaseURL.endsWith("/relax") ||
normalizedBaseURL.endsWith("/mj-relax") ||
normalizedModelId.includes("relax")
) {
return "relax" as const;
}
return "fast" as const;
}
function buildEndpointCandidates(baseURL: string, mode: "fast" | "relax", suffix: string) {
const normalizedBaseURL = normalizeBaseURL(baseURL);
const normalizedSuffix = suffix.startsWith("/") ? suffix : `/${suffix}`;
const directPattern = /\/(fast|relax|mj-fast|mj-relax)$/i;
if (directPattern.test(normalizedBaseURL)) {
return [`${normalizedBaseURL}${normalizedSuffix}`];
}
const modePrefix = mode === "relax" ? "relax" : "fast";
const altPrefix = mode === "relax" ? "mj-relax" : "mj-fast";
return [
`${normalizedBaseURL}/${modePrefix}${normalizedSuffix}`,
`${normalizedBaseURL}/${altPrefix}${normalizedSuffix}`,
];
}
function normalizeImageSizeLabel(imageSize?: string): "1K" | "2K" | "4K" {
const normalized = String(imageSize || "").trim().toUpperCase();
if (normalized === "1K" || normalized === "2K" || normalized === "4K") return normalized;
return DEFAULT_IMAGE_SIZE;
}
function normalizeAspectRatio(value?: string) {
const normalized = String(value || "").trim();
if (!normalized || normalized.toLowerCase() === "auto") return undefined;
if (/^\d+:\d+$/.test(normalized)) return normalized;
return undefined;
}
function appendAspectRatio(prompt: string, aspectRatio?: string) {
const normalizedPrompt = String(prompt || "").trim();
const normalizedAspectRatio = normalizeAspectRatio(aspectRatio);
if (!normalizedAspectRatio) return normalizedPrompt;
if (/--ar\s+\d+:\d+/i.test(normalizedPrompt)) return normalizedPrompt;
return `${normalizedPrompt} --ar ${normalizedAspectRatio}`.trim();
}
function detectImageExtension(imageUrl: string, contentType?: string | null) {
const type = String(contentType || "").toLowerCase();
if (type.includes("jpeg") || type.includes("jpg")) return "jpg";
if (type.includes("webp")) return "webp";
if (type.includes("png")) return "png";
if (type.includes("gif")) return "gif";
const cleanUrl = imageUrl.split("?")[0]?.toLowerCase() || "";
if (cleanUrl.endsWith(".jpg") || cleanUrl.endsWith(".jpeg")) return "jpg";
if (cleanUrl.endsWith(".webp")) return "webp";
if (cleanUrl.endsWith(".gif")) return "gif";
if (cleanUrl.endsWith(".png")) return "png";
return "png";
}
async function createThumbnail(buffer: Buffer) {
try {
return await sharp(buffer)
.resize(400, 400, { fit: "inside", withoutEnlargement: true })
.jpeg({ quality: 80 })
.toBuffer();
} catch (error) {
console.error("[Midjourney] Failed to create thumbnail:", error);
return buffer;
}
}
function generateFilename(prefix: string, extension: string) {
const timestamp = Date.now();
const random = Math.random().toString(36).slice(2, 8);
return `${prefix}-${timestamp}-${random}.${extension}`;
}
async function uploadGeneratedBuffer(
buffer: Buffer,
extension: string,
mimeType: string
): Promise<UploadedImage> {
const thumbnailBuffer = await createThumbnail(buffer);
const originalFilename = generateFilename("midjourney-orig", extension);
const originalArrayBuffer = buffer.buffer.slice(
buffer.byteOffset,
buffer.byteOffset + buffer.byteLength
) as ArrayBuffer;
const originalUploaded = await uploadFile(originalFilename, originalArrayBuffer);
const thumbFilename = generateFilename("midjourney-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 downloadAndUploadImage(imageUrl: string): Promise<UploadedImage> {
const deferred = await createDeferredUploadedImage(imageUrl, {
providerType: "midjourney",
});
if (deferred) return deferred;
const response = await fetch(imageUrl, { cache: "no-store" });
if (!response.ok) {
throw new Error(`Failed to download Midjourney image: ${response.status}`);
}
const buffer = Buffer.from(await response.arrayBuffer());
const extension = detectImageExtension(imageUrl, response.headers.get("content-type"));
const mimeType =
response.headers.get("content-type") ||
`image/${extension === "jpg" ? "jpeg" : extension}`;
return uploadGeneratedBuffer(buffer, extension, mimeType);
}
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 = (value?: string) => {
if (!value) return;
const normalized = String(value).trim();
if (!normalized) 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 };
}
async function loadInputImageAsDataUrl(source: string) {
const normalized = String(source || "").trim();
if (!normalized) {
throw new Error("Midjourney 输入图片地址为空");
}
if (normalized.startsWith("data:")) {
return normalized;
}
const response = await fetch(normalized, { cache: "no-store" });
if (!response.ok) {
throw new Error(`Failed to fetch Midjourney input image: ${response.status}`);
}
const mimeType = response.headers.get("content-type") || "image/png";
const buffer = Buffer.from(await response.arrayBuffer());
return `data:${mimeType};base64,${buffer.toString("base64")}`;
}
async function parseJsonResponse(response: Response) {
const rawText = await response.text();
try {
return JSON.parse(rawText);
} catch {
return rawText;
}
}
function extractErrorMessage(payload: any) {
if (!payload) return "Midjourney 请求失败";
if (typeof payload === "string") return payload;
if (payload.failReason) return String(payload.failReason);
if (payload.description) return String(payload.description);
if (payload.message) return String(payload.message);
if (payload.error?.message) return String(payload.error.message);
try {
return JSON.stringify(payload);
} catch {
return "Midjourney 请求失败";
}
}
async function requestWithEndpointFallback(params: {
baseURL: string;
mode: "fast" | "relax";
suffix: string;
method: "GET" | "POST";
apiKey: string;
body?: string;
contentType?: string;
}) {
const { baseURL, mode, suffix, method, apiKey, body, contentType } = params;
const endpoints = buildEndpointCandidates(baseURL, mode, suffix);
let lastPayload: any = null;
let lastStatus: number | null = null;
let lastEndpoint = endpoints[0];
for (const endpoint of endpoints) {
lastEndpoint = endpoint;
const response = await fetch(endpoint, {
method,
headers: {
Authorization: `Bearer ${apiKey}`,
...(contentType ? { "Content-Type": contentType } : {}),
},
...(body ? { body } : {}),
signal: AbortSignal.timeout(45000),
});
const payload = await parseJsonResponse(response);
lastPayload = payload;
lastStatus = response.status;
if (response.ok || response.status !== 404) {
return { endpoint, response, payload };
}
}
throw new Error(
`Midjourney request failed: ${extractErrorMessage(lastPayload)} (${lastStatus || "no_status"}) @ ${lastEndpoint}`
);
}
async function submitImagineTask(params: {
baseURL: string;
apiKey: string;
mode: "fast" | "relax";
prompt: string;
imageUrls: string[];
}) {
const { baseURL, apiKey, mode, prompt, imageUrls } = params;
const base64Array = await Promise.all(imageUrls.map((imageUrl) => loadInputImageAsDataUrl(imageUrl)));
const { payload } = await requestWithEndpointFallback({
baseURL,
mode,
suffix: "/mj/submit/imagine",
method: "POST",
apiKey,
body: JSON.stringify({
base64Array,
prompt,
}),
contentType: "application/json",
});
const code = Number((payload as any)?.code ?? 0);
const taskId = String((payload as any)?.result || "").trim();
if (![1, 22].includes(code) || !taskId) {
throw new Error(extractErrorMessage(payload));
}
return taskId;
}
async function fetchMidjourneyTask(params: {
baseURL: string;
apiKey: string;
mode: "fast" | "relax";
taskId: string;
}) {
const { payload } = await requestWithEndpointFallback({
baseURL: params.baseURL,
mode: params.mode,
suffix: `/mj/task/${params.taskId}/fetch`,
method: "GET",
apiKey: params.apiKey,
});
return (payload || {}) as MidjourneyTaskResult;
}
async function waitForMidjourneyTask(params: {
baseURL: string;
apiKey: string;
mode: "fast" | "relax";
taskId: string;
}) {
const deadlineAt = Date.now() + DEFAULT_TIMEOUT_MS;
let lastTask: MidjourneyTaskResult | null = null;
while (Date.now() < deadlineAt) {
const task = await fetchMidjourneyTask(params);
lastTask = task;
const status = String(task.status || "").toUpperCase();
if (status === "SUCCESS" && task.imageUrl) {
return task;
}
if (status === "FAILURE") {
throw new Error(task.failReason || task.description || "Midjourney 任务失败");
}
await new Promise((resolve) => setTimeout(resolve, POLL_INTERVAL_MS));
}
throw new Error(
`Midjourney task timeout: ${params.taskId} (${lastTask?.status || "unknown"})`
);
}
function formatMarkdown(image: UploadedImage, imageSize: string) {
return `![thumb|||${image.original}|||${imageSize}](${image.thumbnail})`;
}
export function createMidjourney({ apiKey, baseURL }: MidjourneyConfig) {
return {
provider: "midjourney",
chat: (modelId: string) => ({
provider: "midjourney",
doGenerate: async (options: {
messages: Array<{ role: string; content: any }>;
imageParams?: MidjourneyImageParams;
overrideModelId?: string;
}) => {
const { prompt, imageUrls } = extractPromptAndImages(options.messages);
const activeModelId = options.overrideModelId || modelId;
const mode = resolveMode(baseURL, activeModelId);
const imageSize = normalizeImageSizeLabel(options.imageParams?.imageSize);
const finalPrompt = appendAspectRatio(prompt, options.imageParams?.aspectRatio);
console.log("[Midjourney] Submitting request", {
model: activeModelId,
mode,
imageCount: imageUrls.length,
promptPreview: finalPrompt.slice(0, 120),
});
const taskId = await submitImagineTask({
baseURL,
apiKey,
mode,
prompt: finalPrompt,
imageUrls,
});
console.log("[Midjourney] Task submitted", { taskId, mode });
const taskResult = await waitForMidjourneyTask({
baseURL,
apiKey,
mode,
taskId,
});
if (!taskResult.imageUrl) {
throw new Error(taskResult.failReason || taskResult.description || "Midjourney 未返回图片地址");
}
const uploadedImage = await downloadAndUploadImage(taskResult.imageUrl);
const text = formatMarkdown(uploadedImage, imageSize);
return {
text,
finishReason: "stop" as const,
usage: {
promptTokens: 0,
completionTokens: text.length,
},
providerMetadata: {
midjourney: {
taskId,
mode,
progress: taskResult.progress || "",
buttons: taskResult.buttons || [],
imageUrl: taskResult.imageUrl,
},
},
};
},
doStream: async function* (options: {
messages: Array<{ role: string; content: any }>;
imageParams?: MidjourneyImageParams;
overrideModelId?: string;
}) {
const result = await this.doGenerate(options);
yield {
type: "text-delta",
textDelta: result.text,
};
},
}),
};
}