Initial commit
This commit is contained in:
Jordan
213
mcp-server/tools/media/generateImage.js
Normal file
213
mcp-server/tools/media/generateImage.js
Normal file
@@ -0,0 +1,213 @@
|
||||
import { z } from "zod";
|
||||
import axios from "axios";
|
||||
import fs from "fs";
|
||||
import path from "path";
|
||||
import sharp from "sharp";
|
||||
import { withAuth, getSessionCredentials } from "../../auth/index.js";
|
||||
import { handleToolError } from "../helpers/errorHandler.js";
|
||||
import { withAuthParams } from "../helpers/authSchema.js";
|
||||
|
||||
// --- Verificación de créditos y reporte de uso ---
|
||||
const WS_BASE = "https://ws.cocosolution.com/api/handler_acaicode.php";
|
||||
// Precios Gemini 2.5 Flash: input $0.15/1M tokens, output $0.60/1M tokens
|
||||
function calcCost(usageMetadata) {
|
||||
const input = usageMetadata?.promptTokenCount || 0;
|
||||
const output = usageMetadata?.candidatesTokenCount || 0;
|
||||
return Math.round(((input * 0.15 + output * 0.60) / 1_000_000) * 1e6) / 1e6;
|
||||
}
|
||||
|
||||
function getAcaiToken() {
|
||||
const projectDir = process.env.ACAI_PROJECT_DIR || "";
|
||||
if (!projectDir) return null;
|
||||
try {
|
||||
const acaiFile = path.join(projectDir, ".acai");
|
||||
const data = JSON.parse(fs.readFileSync(acaiFile, "utf-8"));
|
||||
return data.token || null;
|
||||
} catch { return null; }
|
||||
}
|
||||
|
||||
async function checkCredits() {
|
||||
const token = getAcaiToken();
|
||||
if (!token) return false; // Si no hay token, no bloquear
|
||||
const testParam = process.env.STRIPE_MODE === "test" ? "&test" : "";
|
||||
try {
|
||||
const resp = await axios.put(`${WS_BASE}?action=getUsageLimits${testParam}`, {}, {
|
||||
headers: { Authorization: `Bearer ${token}`, "Content-Type": "application/json" },
|
||||
timeout: 10000,
|
||||
});
|
||||
return resp.data?.data?.exceeded === true;
|
||||
} catch { return false; }
|
||||
}
|
||||
|
||||
function reportImageUsage(usageMetadata, model) {
|
||||
const token = getAcaiToken();
|
||||
if (!token) return;
|
||||
const testParam = process.env.STRIPE_MODE === "test" ? "&test" : "";
|
||||
const cost = calcCost(usageMetadata);
|
||||
const payload = {
|
||||
action: "reportUsage",
|
||||
model: model || "gemini-2.5-flash-image",
|
||||
input_tokens: usageMetadata?.promptTokenCount || 0,
|
||||
output_tokens: usageMetadata?.candidatesTokenCount || 0,
|
||||
cache_read_tokens: 0,
|
||||
cache_creation_tokens: 0,
|
||||
cost_usd: cost,
|
||||
session_id: "",
|
||||
};
|
||||
// Fire and forget
|
||||
axios.put(`${WS_BASE}?action=reportUsage${testParam}`, payload, {
|
||||
headers: { Authorization: `Bearer ${token}`, "Content-Type": "application/json" },
|
||||
timeout: 10000,
|
||||
}).then(resp => {
|
||||
if (resp.data?.success) console.error(`[generate_image] Usage reported: ${model} cost=$${cost}`);
|
||||
else console.error(`[generate_image] Usage report failed:`, resp.data);
|
||||
}).catch(err => console.error(`[generate_image] Usage report error:`, err.message));
|
||||
}
|
||||
|
||||
export function registerGenerateImageTool(server) {
|
||||
server.tool(
|
||||
"generate_image",
|
||||
`Generate an AI image and save it to the project's uploads folder. Returns preview URLs plus the recommended upload URL for upload_record_image. In Forge environments, prefer uploadUrl (or fullUrl if uploadUrl is absent) over dockerUrl when assigning the image to a record field.`,
|
||||
withAuthParams({
|
||||
prompt: z.string().describe("Description of the image to generate"),
|
||||
width: z.number().optional().describe("Image width in pixels (default: 1024)"),
|
||||
height: z.number().optional().describe("Image height in pixels (default: 1024)"),
|
||||
style: z.string().optional().describe("Image style hint to add to prompt (e.g., 'photographic', 'digital-art', 'minimalist')"),
|
||||
fileName: z.string().optional().describe("Custom filename (without extension). If not provided, auto-generated."),
|
||||
}),
|
||||
{ readOnlyHint: false, destructiveHint: false },
|
||||
withAuth(async ({ prompt, width = 1024, height = 1024, style, fileName }, extra) => {
|
||||
try {
|
||||
const nanoBananaApiKey = process.env.NANO_BANANA_API_KEY;
|
||||
if (!nanoBananaApiKey) {
|
||||
return {
|
||||
content: [{ type: "text", text: "Error: NANO_BANANA_API_KEY not set." }],
|
||||
isError: true,
|
||||
};
|
||||
}
|
||||
|
||||
const projectDir = process.env.ACAI_PROJECT_DIR || "";
|
||||
if (!projectDir) {
|
||||
return {
|
||||
content: [{ type: "text", text: "Error: ACAI_PROJECT_DIR not set." }],
|
||||
isError: true,
|
||||
};
|
||||
}
|
||||
|
||||
// Verificar créditos antes de generar
|
||||
const exceeded = await checkCredits();
|
||||
if (exceeded) {
|
||||
return {
|
||||
content: [{ type: "text", text: "Error: No te quedan créditos. Mejora tu plan para seguir usando el asistente." }],
|
||||
isError: true,
|
||||
};
|
||||
}
|
||||
|
||||
// Build prompt with style hint
|
||||
const fullPrompt = style ? `${prompt}. Style: ${style}` : prompt;
|
||||
|
||||
// Generate image via Google Gemini
|
||||
const geminiModel = process.env.NANO_BANANA_MODEL || "gemini-2.5-flash-image";
|
||||
const apiUrl = process.env.NANO_BANANA_URL ||
|
||||
`https://generativelanguage.googleapis.com/v1beta/models/${geminiModel}:generateContent`;
|
||||
|
||||
const generateResponse = await axios.post(
|
||||
apiUrl,
|
||||
{
|
||||
contents: [{ parts: [{ text: fullPrompt }] }],
|
||||
generationConfig: {
|
||||
responseModalities: ["TEXT", "IMAGE"],
|
||||
},
|
||||
},
|
||||
{
|
||||
headers: {
|
||||
"x-goog-api-key": nanoBananaApiKey,
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
timeout: 120000,
|
||||
validateStatus: (status) => status < 500,
|
||||
}
|
||||
);
|
||||
|
||||
// Extract image from response
|
||||
let imageBuffer = null;
|
||||
if (generateResponse.data.candidates?.[0]?.content?.parts) {
|
||||
for (const part of generateResponse.data.candidates[0].content.parts) {
|
||||
if (part.inlineData?.data) {
|
||||
imageBuffer = Buffer.from(part.inlineData.data, "base64");
|
||||
break;
|
||||
}
|
||||
if (part.text?.startsWith("data:image")) {
|
||||
const match = part.text.match(/data:image\/[^;]+;base64,(.+)/);
|
||||
if (match) {
|
||||
imageBuffer = Buffer.from(match[1], "base64");
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (!imageBuffer) {
|
||||
return {
|
||||
content: [{
|
||||
type: "text",
|
||||
text: `Error: Could not extract image from API response. Status: ${generateResponse.status}. Response: ${JSON.stringify(generateResponse.data).substring(0, 1000)}`
|
||||
}],
|
||||
isError: true,
|
||||
};
|
||||
}
|
||||
|
||||
// Compress to JPEG
|
||||
const originalSize = imageBuffer.length;
|
||||
try {
|
||||
imageBuffer = await sharp(imageBuffer)
|
||||
.jpeg({ quality: 85 })
|
||||
.toBuffer();
|
||||
console.error(`[generate_image] Compressed: ${Math.round(originalSize / 1024)}KB → ${Math.round(imageBuffer.length / 1024)}KB`);
|
||||
} catch (e) {
|
||||
console.error(`[generate_image] Compression failed, using original:`, e.message);
|
||||
}
|
||||
|
||||
// Save to cms/uploads/generated/
|
||||
const uploadsDir = path.join(projectDir, "cms", "uploads", "generated");
|
||||
fs.mkdirSync(uploadsDir, { recursive: true });
|
||||
|
||||
const safeName = fileName
|
||||
? fileName.replace(/[^\w\-]/g, "_") + ".jpg"
|
||||
: `generated-${Date.now()}.jpg`;
|
||||
const filePath = path.join(uploadsDir, safeName);
|
||||
fs.writeFileSync(filePath, imageBuffer);
|
||||
|
||||
const relativePath = `cms/uploads/generated/${safeName}`;
|
||||
const dockerUrl = `http://localhost/${relativePath}`;
|
||||
|
||||
// Reportar uso (fire and forget)
|
||||
reportImageUsage(generateResponse.data.usageMetadata, geminiModel);
|
||||
|
||||
const credentials = await getSessionCredentials(extra.sessionId);
|
||||
const fullUrl = credentials.web_url ? `${credentials.web_url}/${relativePath}` : dockerUrl;
|
||||
const uploadUrl = credentials.web_url ? fullUrl : dockerUrl;
|
||||
|
||||
return {
|
||||
content: [{
|
||||
type: "text",
|
||||
text: JSON.stringify({
|
||||
success: true,
|
||||
prompt: fullPrompt,
|
||||
fileName: safeName,
|
||||
filePath,
|
||||
relativePath,
|
||||
dockerUrl,
|
||||
fullUrl,
|
||||
uploadUrl,
|
||||
size: `${Math.round(imageBuffer.length / 1024)}KB`,
|
||||
note: `Image saved. To assign it with upload_record_image, use imageUrl="${uploadUrl}". dockerUrl is mainly for local preview/debugging.`,
|
||||
}, null, 2),
|
||||
}],
|
||||
};
|
||||
} catch (error) {
|
||||
return handleToolError(error, "generate_image", { prompt });
|
||||
}
|
||||
})
|
||||
);
|
||||
}
|
||||
Reference in New Issue
Block a user