tools.js•7.75 kB
import axios from "axios";
import { CallToolRequestSchema, ListToolsRequestSchema, } from "@modelcontextprotocol/sdk/types.js";
import open from "open";
import { createImageTool } from "./tools/createImage.js";
import { upscaleImageTool } from "./tools/upscaleImage.js";
export const setTools = (server) => {
server.setRequestHandler(ListToolsRequestSchema, async () => {
return {
tools: [createImageTool, upscaleImageTool],
};
});
server.setRequestHandler(CallToolRequestSchema, async (request, connection) => {
if (request.params.name === "letzai_create_image") {
try {
let { prompt, width, height, quality, creativity, hasWatermark, systemVersion, mode, } = request.params.arguments;
mode = !mode || !mode.includes(mode || "") ? "turbo" : mode;
width = parseInt(width) || 1600;
height = parseInt(height) || 1600;
quality = parseInt(quality) || 2;
creativity = parseInt(creativity) || 2;
systemVersion = parseInt(systemVersion) || 3;
hasWatermark =
typeof hasWatermark === "boolean" ? hasWatermark : false;
// Step 1: Create the image request
const responseCreate = await axios.post("https://api.letz.ai/images", {
prompt,
width,
height,
quality,
creativity,
hasWatermark,
systemVersion,
mode,
}, {
headers: {
Authorization: `Bearer ${process.env.LETZAI_API_KEY}`,
},
});
let imageFinished = false;
let imageVersions = null;
let imageId = responseCreate.data.id;
// Step 2: Poll for image creation status
while (!imageFinished) {
await new Promise((resolve) => setTimeout(resolve, 5000)); // Wait before checking again
const responseImage = await axios.get(`https://api.letz.ai/images/${imageId}`, {
headers: {
Authorization: `Bearer ${process.env.LETZAI_API_KEY}`,
},
});
if (responseImage.data.progress < 100) {
// Send a progress notification (through stdout for Stdio transport)
console.log(JSON.stringify({
jsonrpc: "2.0",
method: "progress_update",
params: {
message: `Image is still being processed. Progress: ${responseImage.data.progress}%`,
},
}));
}
else {
imageFinished = true;
imageVersions = responseImage.data.imageVersions;
}
}
// Convert the image to Base64 after processing is complete
/* const imageBase64 = convertImageUrlToBase64(
imageVersions?.["640x640"] as string
);
*/
// Open the image in browser
open(imageVersions?.original);
// Return the response to the client
return {
content: [
{
type: "text",
text: `Image generated successfully!\nThe image has been opened in your default browser.\n\n Image URL: ${imageVersions?.original}\n\nYou can also click the URL above to view the image again.`,
},
],
};
}
catch (err) {
return {
content: [
{
type: "text",
text: `Error happened: ${err.toString()}`,
},
],
};
}
}
else if (request.params.name === "letzai_upscale_image") {
try {
let { imageId, imageUrl, strength } = request.params.arguments;
strength = parseInt(strength) || 1;
let body = {};
if (imageId) {
body = {
imageId,
strength,
};
}
else if (imageUrl) {
body = {
imageUrl,
strength,
};
}
else {
throw new Error("Provide image ID or Image URL");
}
// Step 1: Create the image request
const responseCreate = await axios.post("https://api.letz.ai/upscale", body, {
headers: {
Authorization: `Bearer ${process.env.LETZAI_API_KEY}`,
},
});
let imageFinished = false;
let imageVersions = null;
let upscaleId = responseCreate.data.id;
// Step 2: Poll for image creation status
while (!imageFinished) {
await new Promise((resolve) => setTimeout(resolve, 5000)); // Wait before checking again
const responseImage = await axios.get(`https://api.letz.ai/upscale/${upscaleId}`, {
headers: {
Authorization: `Bearer ${process.env.LETZAI_API_KEY}`,
},
});
if (responseImage.data.status != "ready") {
// Send a progress notification (through stdout for Stdio transport)
console.log(JSON.stringify({
jsonrpc: "2.0",
method: "progress_update",
params: {
message: `Image is still being processed. Progress: ${responseImage.data.progress}%`,
},
}));
}
else {
imageFinished = true;
imageVersions = responseImage.data.imageVersions;
}
}
// Convert the image to Base64 after processing is complete
/* const imageBase64 = convertImageUrlToBase64(
imageVersions?.["640x640"] as string
);
*/
// Open the image in browser
open(imageVersions?.original);
// Return the response to the client
return {
content: [
{
type: "text",
text: `Image upscaled successfully!\nThe image has been opened in your default browser.\n\n Image URL: ${imageVersions?.original}\n\nYou can also click the URL above to view the image again.`,
},
],
};
}
catch (err) {
return {
content: [
{
type: "text",
text: `Error happened: ${err.toString()}`,
},
],
};
}
}
throw new Error("Tool not found");
});
};