218 lines
10 KiB
TypeScript
218 lines
10 KiB
TypeScript
import sharp from "sharp";
|
|
|
|
import type {
|
|
GenerationAdapter,
|
|
GenerationAdapterRequest,
|
|
GenerationAdapterResult,
|
|
NormalizedGenerationOutput,
|
|
} from "./ai-adapter-contract.js";
|
|
import { gptImageRequestSizeForRatio, normalizeImageOutputToRatio } from "./image-output-normalizer.mjs";
|
|
|
|
const geminiProductModelId = "gemini-3.1-flash-image-preview";
|
|
const geminiProviderModelId = "gemini-3.1-flash-image";
|
|
const gptImageModelId = "gpt-image-2";
|
|
const geminiEndpoint = "https://oneapi.intelligrow.cn/v1/chat/completions";
|
|
const gptImageEndpoint = "https://oneapi.intelligrow.cn/v1/images/generations";
|
|
const gptImageReferenceEndpoint = "https://oneapi.intelligrow.cn/v1/images/edits";
|
|
const maximumResponseBytes = 32 * 1024 * 1024;
|
|
const requestTimeoutMilliseconds = 180_000;
|
|
|
|
type FetchLike = typeof fetch;
|
|
|
|
class OneApiRuntimeError extends Error {
|
|
constructor(
|
|
readonly category: "gateway_balance_insufficient" | "gateway_contract_invalid" | "model_disabled" | "reference_invalid" | "upstream_failed" | "upstream_timeout" | "unknown_non_retryable",
|
|
readonly sourceCategory: string,
|
|
) {
|
|
super(sourceCategory);
|
|
}
|
|
}
|
|
|
|
function failure(error: unknown): GenerationAdapterResult {
|
|
if (error instanceof OneApiRuntimeError) {
|
|
return { category: error.category, sourceCategory: error.sourceCategory, status: "failed" };
|
|
}
|
|
if (error instanceof Error && error.name === "AbortError") {
|
|
return { category: "upstream_timeout", sourceCategory: "upstream_timeout", status: "failed" };
|
|
}
|
|
return { category: "upstream_failed", sourceCategory: "upstream_failed", status: "failed" };
|
|
}
|
|
|
|
function mapHttpFailure(status: number) {
|
|
if (status === 408 || status === 504) return new OneApiRuntimeError("upstream_timeout", `upstream_http_${status}`);
|
|
if (status === 429) return new OneApiRuntimeError("gateway_balance_insufficient", "upstream_http_429");
|
|
if (status >= 500) return new OneApiRuntimeError("upstream_failed", `upstream_http_${status}`);
|
|
if (status === 400 || status === 404 || status === 422) return new OneApiRuntimeError("gateway_contract_invalid", `upstream_http_${status}`);
|
|
return new OneApiRuntimeError("unknown_non_retryable", `upstream_http_${status}`);
|
|
}
|
|
|
|
async function readBoundedJson(response: Response) {
|
|
const declaredLength = Number(response.headers.get("content-length") ?? 0);
|
|
if (Number.isFinite(declaredLength) && declaredLength > maximumResponseBytes) {
|
|
throw new OneApiRuntimeError("gateway_contract_invalid", "upstream_response_too_large");
|
|
}
|
|
if (!response.body) throw new OneApiRuntimeError("gateway_contract_invalid", "upstream_response_empty");
|
|
const reader = response.body.getReader();
|
|
const chunks: Buffer[] = [];
|
|
let total = 0;
|
|
try {
|
|
while (true) {
|
|
const next = await reader.read();
|
|
if (next.done) break;
|
|
const chunk = Buffer.from(next.value);
|
|
total += chunk.length;
|
|
if (total > maximumResponseBytes) {
|
|
await reader.cancel();
|
|
throw new OneApiRuntimeError("gateway_contract_invalid", "upstream_response_too_large");
|
|
}
|
|
chunks.push(chunk);
|
|
}
|
|
try {
|
|
return JSON.parse(Buffer.concat(chunks).toString("utf8")) as unknown;
|
|
} catch {
|
|
throw new OneApiRuntimeError("gateway_contract_invalid", "upstream_response_invalid");
|
|
}
|
|
} finally {
|
|
for (const chunk of chunks) chunk.fill(0);
|
|
}
|
|
}
|
|
|
|
function extractGeminiImage(response: unknown) {
|
|
if (!response || typeof response !== "object" || !("choices" in response) || !Array.isArray(response.choices)) {
|
|
throw new OneApiRuntimeError("gateway_contract_invalid", "response_shape_invalid");
|
|
}
|
|
const choice = response.choices[0];
|
|
const content = choice && typeof choice === "object" && "message" in choice && choice.message && typeof choice.message === "object"
|
|
&& "content" in choice.message && typeof choice.message.content === "string" ? choice.message.content : "";
|
|
const matches = [...content.matchAll(/!\[[^\]]*\]\(\s*data:(image\/(?:jpeg|png|webp));base64,([A-Za-z0-9+/=\r\n]+)\s*\)/gi)];
|
|
if (matches.length !== 1) throw new OneApiRuntimeError("gateway_contract_invalid", "response_single_image_required");
|
|
return { bytes: Buffer.from(matches[0]![2]!, "base64"), declaredMimeType: matches[0]![1]!.toLowerCase() };
|
|
}
|
|
|
|
function extractGptImage(response: unknown) {
|
|
if (!response || typeof response !== "object" || !("data" in response) || !Array.isArray(response.data)
|
|
|| response.data.length !== 1 || !response.data[0] || typeof response.data[0] !== "object"
|
|
|| !("b64_json" in response.data[0]) || typeof response.data[0].b64_json !== "string") {
|
|
throw new OneApiRuntimeError("gateway_contract_invalid", "response_single_image_required");
|
|
}
|
|
return { bytes: Buffer.from(response.data[0].b64_json, "base64"), declaredMimeType: undefined };
|
|
}
|
|
|
|
async function normalizeOutput(bytes: Buffer, declaredMimeType: string | undefined, ratio: GenerationAdapterRequest["ratio"]): Promise<NormalizedGenerationOutput> {
|
|
try {
|
|
const metadata = await sharp(bytes, { failOn: "error", limitInputPixels: 40_000_000 }).metadata();
|
|
const mimeType = metadata.format === "png" ? "image/png" : metadata.format === "jpeg" ? "image/jpeg" : metadata.format === "webp" ? "image/webp" : undefined;
|
|
if (!mimeType || !metadata.width || !metadata.height || (declaredMimeType && declaredMimeType !== mimeType)) {
|
|
throw new OneApiRuntimeError("gateway_contract_invalid", "response_media_invalid");
|
|
}
|
|
const normalized = await normalizeImageOutputToRatio({ bytes, mimeType, pixelHeight: metadata.height, pixelWidth: metadata.width, ratio });
|
|
return { bytes: normalized.bytes, mimeType: normalized.mimeType, pixelHeight: normalized.pixelHeight, pixelWidth: normalized.pixelWidth };
|
|
} catch (error) {
|
|
if (error instanceof OneApiRuntimeError) throw error;
|
|
throw new OneApiRuntimeError("gateway_contract_invalid", "response_media_invalid");
|
|
}
|
|
}
|
|
|
|
function validateRequest(request: GenerationAdapterRequest) {
|
|
if (!request.prompt.trim() || request.prompt.length > 1_000) throw new OneApiRuntimeError("gateway_contract_invalid", "prompt_invalid");
|
|
const references = request.referenceImages ?? [];
|
|
if (references.length !== request.referenceAssetIds.length || references.length > 2) {
|
|
throw new OneApiRuntimeError("reference_invalid", "reference_count_invalid");
|
|
}
|
|
const totalBytes = references.reduce((total, reference) => total + reference.bytes.length, 0);
|
|
if (totalBytes > 20 * 1024 * 1024 || references.some((reference) => reference.bytes.length === 0 || reference.bytes.length > 10 * 1024 * 1024)) {
|
|
throw new OneApiRuntimeError("reference_invalid", "reference_size_invalid");
|
|
}
|
|
return references;
|
|
}
|
|
|
|
function buildRequest(request: GenerationAdapterRequest) {
|
|
const references = validateRequest(request);
|
|
if (request.modelId === geminiProductModelId) {
|
|
const content = references.length === 0
|
|
? request.prompt
|
|
: [
|
|
{ text: request.prompt, type: "text" },
|
|
...references.map((reference) => ({
|
|
image_url: { url: `data:${reference.mimeType};base64,${reference.bytes.toString("base64")}` },
|
|
type: "image_url",
|
|
})),
|
|
];
|
|
return {
|
|
body: JSON.stringify({
|
|
extra_body: { google: { image_config: { aspect_ratio: request.ratio, image_size: "1K" } } },
|
|
messages: [{ content, role: "user" }],
|
|
model: geminiProviderModelId,
|
|
stream: false,
|
|
}),
|
|
contentType: "application/json",
|
|
endpoint: geminiEndpoint,
|
|
parser: extractGeminiImage,
|
|
};
|
|
}
|
|
if (request.modelId !== gptImageModelId) throw new OneApiRuntimeError("model_disabled", "model_not_supported");
|
|
if (references.length > 0) {
|
|
const form = new FormData();
|
|
form.append("model", gptImageModelId);
|
|
form.append("prompt", request.prompt);
|
|
form.append("response_format", "b64_json");
|
|
form.append("size", gptImageRequestSizeForRatio(request.ratio));
|
|
references.forEach((reference, index) => form.append("image[]", new Blob([reference.bytes], { type: reference.mimeType }), `reference-${index + 1}.png`));
|
|
return { body: form, contentType: undefined, endpoint: gptImageReferenceEndpoint, parser: extractGptImage };
|
|
}
|
|
return {
|
|
body: JSON.stringify({ model: gptImageModelId, prompt: request.prompt, response_format: "b64_json", size: gptImageRequestSizeForRatio(request.ratio) }),
|
|
contentType: "application/json",
|
|
endpoint: gptImageEndpoint,
|
|
parser: extractGptImage,
|
|
};
|
|
}
|
|
|
|
export class OneApiGenerationAdapter implements GenerationAdapter {
|
|
private readonly credential: Buffer;
|
|
private readonly fetchImpl: FetchLike;
|
|
private disposed = false;
|
|
|
|
constructor(input: { credential: Buffer; fetch?: FetchLike }) {
|
|
if (input.credential.length < 8) throw new Error("ai_gateway_credential_invalid");
|
|
this.credential = Buffer.from(input.credential);
|
|
this.fetchImpl = input.fetch ?? fetch;
|
|
}
|
|
|
|
async start(request: GenerationAdapterRequest): Promise<GenerationAdapterResult> {
|
|
if (this.disposed) return { category: "upstream_failed", sourceCategory: "adapter_disposed", status: "failed" };
|
|
const controller = new AbortController();
|
|
const timeout = setTimeout(() => controller.abort(), requestTimeoutMilliseconds);
|
|
let sourceBytes: Buffer | undefined;
|
|
try {
|
|
const providerRequest = buildRequest(request);
|
|
const headers = new Headers({ authorization: `Bearer ${this.credential.toString("utf8")}` });
|
|
if (providerRequest.contentType) headers.set("content-type", providerRequest.contentType);
|
|
const response = await this.fetchImpl(providerRequest.endpoint, {
|
|
body: providerRequest.body,
|
|
headers,
|
|
method: "POST",
|
|
redirect: "error",
|
|
signal: controller.signal,
|
|
});
|
|
if (!response.ok) throw mapHttpFailure(response.status);
|
|
const parsed = await readBoundedJson(response);
|
|
const extracted = providerRequest.parser(parsed);
|
|
sourceBytes = extracted.bytes;
|
|
const output = await normalizeOutput(sourceBytes, extracted.declaredMimeType, request.ratio);
|
|
return { outputs: [output], status: "completed" };
|
|
} catch (error) {
|
|
return failure(error);
|
|
} finally {
|
|
clearTimeout(timeout);
|
|
sourceBytes?.fill(0);
|
|
}
|
|
}
|
|
|
|
dispose() {
|
|
if (this.disposed) return;
|
|
this.disposed = true;
|
|
this.credential.fill(0);
|
|
}
|
|
}
|