ima2-gen 1.1.13 → 1.1.16

package/lib/grokImageAdapter.js

@@ -0,0 +1,364 @@
+import { logEvent } from "./logger.js";
+import { mapSizeToGrokImageParams } from "./grokSizeMapper.js";
+import { detectImageMimeFromB64 } from "./refs.js";
+import { getGrokProxyUrl } from "./grokRuntime.js";
+function getGrokEndpoint(ctx, path = "/v1/images/generations") {
+    return {
+        url: getGrokProxyUrl(ctx, path),
+        headers: { "Content-Type": "application/json", Authorization: "Bearer dummy" },
+    };
+}
+function getGrokTimeout(ctx) {
+    return ctx.config.grokProvider?.generationTimeoutMs || 120_000;
+}
+export function grokError(message, status, code) {
+    const err = new Error(message);
+    err.status = status;
+    err.code = code;
+    return err;
+}
+function grokStageError(stage, message, status) {
+    const prefix = stage === "search" ? "GROK_SEARCH" : "GROK_PLANNER";
+    if (status === 429)
+        return grokError(`${stage} rate limited: ${message}`, 429, "GROK_RATE_LIMITED");
+    if (status === 401 || status === 403)
+        return grokError(`${stage} auth failed: ${message}`, 502, "GROK_AUTH_FAILED");
+    if (status >= 500)
+        return grokError(`${stage} upstream error: ${message}`, 502, "GROK_UPSTREAM_ERROR");
+    return grokError(`Grok ${stage} bad request: ${message}`, status, `${prefix}_BAD_REQUEST`);
+}
+function getPlannerConfig(ctx) {
+    const grokCfg = ctx.config.grokProvider || {};
+    return {
+        model: grokCfg.plannerModel || "grok-4.3",
+        timeoutMs: grokCfg.plannerTimeoutMs || 60_000,
+    };
+}
+function withTimeoutSignal(signal, timeoutMs) {
+    const timeoutController = new AbortController();
+    const timer = setTimeout(() => timeoutController.abort(), timeoutMs);
+    const combinedSignal = signal ? AbortSignal.any([signal, timeoutController.signal]) : timeoutController.signal;
+    return { combinedSignal, timer };
+}
+export function imagePayload(model, prompt, size) {
+    return { model, prompt, n: 1, response_format: "b64_json", ...mapSizeToGrokImageParams(size) };
+}
+function referenceImageUrl(ref) {
+    const inputMime = ref.declaredMime || ref.detectedMime || detectImageMimeFromB64(ref.b64) || "image/png";
+    return ref.b64.startsWith("data:") ? ref.b64 : `data:${inputMime};base64,${ref.b64}`;
+}
+export function imageEditPayload(model, prompt, references, size) {
+    const sourceImages = references.map((ref) => ({ type: "image_url", url: referenceImageUrl(ref) }));
+    return { model, prompt, n: 1, response_format: "b64_json", ...(sourceImages.length === 1 ? { image: sourceImages[0] } : { images: sourceImages }), ...mapSizeToGrokImageParams(size) };
+}
+function extractResponsesText(response) {
+    const chunks = [];
+    for (const item of response.output || []) {
+        if (item.type !== "message")
+            continue;
+        for (const content of item.content || []) {
+            if (typeof content.text === "string" && content.text.trim())
+                chunks.push(content.text.trim());
+        }
+    }
+    return chunks.join("\n\n").trim();
+}
+export async function postGrokImages(ctx, payload, signal, path = "/v1/images/generations") {
+    const { url, headers } = getGrokEndpoint(ctx, path);
+    const timeoutMs = getGrokTimeout(ctx);
+    const { combinedSignal, timer } = withTimeoutSignal(signal, timeoutMs);
+    try {
+        const res = await fetch(url, {
+            method: "POST",
+            headers,
+            body: JSON.stringify(payload),
+            signal: combinedSignal,
+        });
+        clearTimeout(timer);
+        if (!res.ok) {
+            const text = await res.text().catch(() => "");
+            let parsed;
+            try {
+                parsed = JSON.parse(text);
+            }
+            catch { /* ignore */ }
+            const msg = parsed?.error || text || `HTTP ${res.status}`;
+            if (res.status === 429)
+                throw grokError(`Grok rate limited: ${msg}`, 429, "GROK_RATE_LIMITED");
+            if (res.status === 401 || res.status === 403)
+                throw grokError(`Grok auth failed: ${msg}`, 502, "GROK_AUTH_FAILED");
+            if (res.status >= 500)
+                throw grokError(`Grok upstream error: ${msg}`, 502, "GROK_UPSTREAM_ERROR");
+            throw grokError(`Grok bad request: ${msg}`, res.status, "GROK_BAD_REQUEST");
+        }
+        return await res.json();
+    }
+    catch (e) {
+        clearTimeout(timer);
+        if (e.name === "AbortError") {
+            if (signal?.aborted)
+                throw grokError("Generation canceled", 499, "GENERATION_CANCELED");
+            throw grokError("Grok image generation timed out", 504, "GENERATION_TIMEOUT");
+        }
+        if (e.code && e.status)
+            throw e;
+        throw grokError(`Grok request failed: ${e.message}`, 502, "GROK_NETWORK_FAILED");
+    }
+}
+export function buildGrokPlannerPayload(prompt, model, size, sizeParams, plannerModel = "grok-4.3", searchSummary = "", references = 0) {
+    const referenceImages = Array.isArray(references) ? references : [];
+    const referenceCount = Array.isArray(references) ? references.length : references;
+    const sizeLine = size
+        ? `Requested ima2 size: ${size}; xAI parameters: ${JSON.stringify(sizeParams)}.`
+        : `Requested ima2 size: auto; xAI parameters: ${JSON.stringify(sizeParams)}.`;
+    const referenceLine = referenceCount > 0
+        ? `Reference images attached: ${referenceCount}. The final image call will use xAI image editing with these source images; preserve relevant subject, composition, style, and product details from them.`
+        : "Reference images attached: 0. The final image call will use text-to-image generation.";
+    return {
+        model: plannerModel,
+        stream: false,
+        parallel_tool_calls: false,
+        messages: [
+            {
+                role: "system",
+                content: [
+                    "You are ima2's image generation planner for xAI Grok Imagine.",
+                    "Rewrite the user's request into one concise, production-ready image prompt.",
+                    "Preserve the user's intent, visible text, subject identity, composition, and style constraints.",
+                    "The final image prompt argument MUST be written in English, even when the user prompt is not English.",
+                    "If the user requests exact visible text in another language, keep that visible text verbatim inside the English prompt.",
+                    "If reference images are attached, inspect them directly and describe the visual details needed for the final image call.",
+                    "Call generate_image exactly once. Do not answer with plain text.",
+                ].join(" "),
+            },
+            {
+                role: "user",
+                content: [
+                    {
+                        type: "text",
+                        text: [
+                            `Selected image model: ${model}.`,
+                            sizeLine,
+                            referenceLine,
+                            searchSummary ? `Mandatory web-search brief:\n${searchSummary}` : "Mandatory web-search brief: unavailable.",
+                            "Create the best final prompt for the image generator.",
+                            "Return the generate_image.prompt argument in English only, except for exact visible text that the user explicitly requested.",
+                            "",
+                            "User prompt:",
+                            prompt,
+                        ].join("\n"),
+                    },
+                    ...referenceImages.map((ref) => ({
+                        type: "image_url",
+                        image_url: { url: referenceImageUrl(ref), detail: "high" },
+                    })),
+                ],
+            },
+        ],
+        tools: [
+            {
+                type: "function",
+                function: {
+                    name: "generate_image",
+                    description: "Generate a single image through xAI Images API.",
+                    parameters: {
+                        type: "object",
+                        properties: {
+                            prompt: {
+                                type: "string",
+                                description: "Final image-generation prompt to send to xAI Images API.",
+                            },
+                            model: {
+                                type: "string",
+                                enum: ["grok-imagine-image", "grok-imagine-image-quality"],
+                                description: "The xAI image model. The server may override this with the user's selected model.",
+                            },
+                        },
+                        required: ["prompt", "model"],
+                    },
+                },
+            },
+        ],
+        tool_choice: { type: "function", function: { name: "generate_image" } },
+    };
+}
+export function buildGrokSearchPayload(prompt, plannerModel = "grok-4.3") {
+    return {
+        model: plannerModel,
+        stream: false,
+        input: [
+            {
+                role: "system",
+                content: [
+                    "You are ima2's visual research assistant.",
+                    "You must use web_search before producing the brief.",
+                    "Return a concise image-generation research brief: visual facts, current references, style cues, and text-rendering constraints.",
+                    "Do not generate an image prompt yet.",
+                ].join(" "),
+            },
+            {
+                role: "user",
+                content: prompt,
+            },
+        ],
+        tools: [{ type: "web_search" }],
+        tool_choice: "required",
+    };
+}
+export async function searchGrokVisualContext(prompt, ctx, options = {}) {
+    const planner = getPlannerConfig(ctx);
+    const payload = buildGrokSearchPayload(prompt, planner.model);
+    const { url, headers } = getGrokEndpoint(ctx, "/v1/responses");
+    const { combinedSignal, timer } = withTimeoutSignal(options.signal, planner.timeoutMs);
+    logEvent("grok", "search:start", { requestId: options.requestId, plannerModel: planner.model, promptChars: prompt.length });
+    try {
+        const res = await fetch(url, {
+            method: "POST",
+            headers,
+            body: JSON.stringify(payload),
+            signal: combinedSignal,
+        });
+        clearTimeout(timer);
+        if (!res.ok) {
+            const text = await res.text().catch(() => "");
+            let parsed;
+            try {
+                parsed = JSON.parse(text);
+            }
+            catch { /* ignore */ }
+            const msg = parsed?.error || text || `HTTP ${res.status}`;
+            throw grokStageError("search", msg, res.status);
+        }
+        const summary = extractResponsesText(await res.json());
+        if (!summary)
+            throw grokError("Grok web search returned no research summary", 502, "GROK_SEARCH_EMPTY_RESPONSE");
+        logEvent("grok", "search:done", { requestId: options.requestId, plannerModel: planner.model, summaryChars: summary.length });
+        return { summary };
+    }
+    catch (e) {
+        clearTimeout(timer);
+        if (e.name === "AbortError") {
+            if (options.signal?.aborted)
+                throw grokError("Generation canceled", 499, "GENERATION_CANCELED");
+            throw grokError("Grok web search timed out", 504, "GROK_SEARCH_TIMEOUT");
+        }
+        if (e.code && e.status)
+            throw e;
+        throw grokError(`Grok web search request failed: ${e.message}`, 502, "GROK_SEARCH_NETWORK_FAILED");
+    }
+}
+export function parseGrokImagePlan(response, fallbackModel) {
+    const toolCalls = response.choices?.[0]?.message?.tool_calls || [];
+    const call = toolCalls.find((item) => item.type === "function" && item.function?.name === "generate_image");
+    if (!call?.function?.arguments) {
+        throw grokError("Grok planner did not call generate_image", 502, "GROK_PLANNER_EMPTY_TOOL_CALL");
+    }
+    let args;
+    try {
+        args = JSON.parse(call.function.arguments);
+    }
+    catch {
+        throw grokError("Grok planner returned invalid tool arguments", 502, "GROK_PLANNER_INVALID_TOOL_ARGS");
+    }
+    if (typeof args?.prompt !== "string" || !args.prompt.trim()) {
+        throw grokError("Grok planner returned an empty image prompt", 502, "GROK_PLANNER_INVALID_TOOL_ARGS");
+    }
+    return { prompt: args.prompt.trim(), model: fallbackModel, webSearchCalls: 1 };
+}
+export async function planGrokImage(prompt, ctx, options = {}) {
+    const imageModel = options.model || ctx.config.grokProvider?.defaultImageModel || "grok-imagine-image";
+    const planner = getPlannerConfig(ctx);
+    const sizeParams = mapSizeToGrokImageParams(options.size);
+    const search = await searchGrokVisualContext(prompt, ctx, { signal: options.signal, requestId: options.requestId });
+    const payload = buildGrokPlannerPayload(prompt, imageModel, options.size, sizeParams, planner.model, search.summary, options.references || options.referenceCount || 0);
+    const { url, headers } = getGrokEndpoint(ctx, "/v1/chat/completions");
+    const { combinedSignal, timer } = withTimeoutSignal(options.signal, planner.timeoutMs);
+    logEvent("grok", "planner:start", { requestId: options.requestId, plannerModel: planner.model, imageModel, size: options.size });
+    try {
+        const res = await fetch(url, {
+            method: "POST",
+            headers,
+            body: JSON.stringify(payload),
+            signal: combinedSignal,
+        });
+        clearTimeout(timer);
+        if (!res.ok) {
+            const text = await res.text().catch(() => "");
+            let parsed;
+            try {
+                parsed = JSON.parse(text);
+            }
+            catch { /* ignore */ }
+            const msg = parsed?.error || text || `HTTP ${res.status}`;
+            throw grokStageError("planner", msg, res.status);
+        }
+        const plan = parseGrokImagePlan(await res.json(), imageModel);
+        logEvent("grok", "planner:done", {
+            requestId: options.requestId,
+            plannerModel: planner.model,
+            imageModel,
+            promptChars: plan.prompt.length,
+            aspectRatio: sizeParams.aspect_ratio,
+            resolution: sizeParams.resolution,
+        });
+        return plan;
+    }
+    catch (e) {
+        clearTimeout(timer);
+        if (e.name === "AbortError") {
+            if (options.signal?.aborted)
+                throw grokError("Generation canceled", 499, "GENERATION_CANCELED");
+            throw grokError("Grok planner timed out", 504, "GROK_PLANNER_TIMEOUT");
+        }
+        if (e.code && e.status)
+            throw e;
+        throw grokError(`Grok planner request failed: ${e.message}`, 502, "GROK_PLANNER_NETWORK_FAILED");
+    }
+}
+export async function generateViaGrok(prompt, ctx, options = {}) {
+    const model = options.model || ctx.config.grokProvider?.defaultImageModel || "grok-imagine-image";
+    const references = options.references || [];
+    const plan = options.plannedPrompt
+        ? { prompt: options.plannedPrompt, model, webSearchCalls: options.webSearchCalls ?? 1 }
+        : await planGrokImage(prompt, ctx, { ...options, referenceCount: references.length });
+    const hasReferences = references.length > 0;
+    const payload = hasReferences
+        ? imageEditPayload(model, plan.prompt, references, options.size)
+        : imagePayload(model, plan.prompt, options.size);
+    const endpoint = hasReferences ? "/v1/images/edits" : "/v1/images/generations";
+    const logStage = hasReferences ? "generate:edit-start" : "generate:start";
+    logEvent("grok", logStage, {
+        requestId: options.requestId,
+        model,
+        promptChars: plan.prompt.length,
+        size: options.size,
+        refs: references.length,
+    });
+    const result = await postGrokImages(ctx, payload, options.signal, endpoint);
+    if (!result.data?.[0]?.b64_json) {
+        throw grokError("Grok returned empty image data", 502, "GROK_EMPTY_RESPONSE");
+    }
+    const usage = result.usage ? { grok_cost_usd_ticks: result.usage.cost_in_usd_ticks ?? 0 } : null;
+    logEvent("grok", "generate:done", {
+        requestId: options.requestId,
+        model,
+        endpoint,
+        refs: references.length,
+        b64Len: result.data[0].b64_json.length,
+    });
+    return { b64: result.data[0].b64_json, usage, webSearchCalls: plan.webSearchCalls, mime: result.data[0].mime_type, revisedPrompt: plan.prompt };
+}
+export async function editViaGrok(prompt, imageB64, ctx, options = {}) {
+    const model = options.model || ctx.config.grokProvider?.defaultImageModel || "grok-imagine-image";
+    const detectedInputMime = detectImageMimeFromB64(imageB64) || "image/png";
+    const imageUrl = imageB64.startsWith("data:") ? imageB64 : `data:${detectedInputMime};base64,${imageB64}`;
+    const payload = { model, prompt, n: 1, response_format: "b64_json", image: { type: "image_url", url: imageUrl }, ...mapSizeToGrokImageParams(options.size) };
+    logEvent("grok", "edit:start", { requestId: options.requestId, model, promptChars: prompt.length });
+    const result = await postGrokImages(ctx, payload, options.signal, "/v1/images/edits");
+    if (!result.data?.[0]?.b64_json) {
+        throw grokError("Grok edit returned empty image data", 502, "GROK_EMPTY_RESPONSE");
+    }
+    const usage = result.usage ? { grok_cost_usd_ticks: result.usage.cost_in_usd_ticks ?? 0 } : null;
+    logEvent("grok", "edit:done", { requestId: options.requestId, model, b64Len: result.data[0].b64_json.length });
+    return { b64: result.data[0].b64_json, usage, webSearchCalls: 0, mime: result.data[0].mime_type, revisedPrompt: result.data[0].revised_prompt || prompt };
+}