pollo-mcp
Version:
Pollo AI Model Context Protocol (MCP) Server for video generation
183 lines (181 loc) • 7.17 kB
JavaScript
var __awaiter = (this && this.__awaiter) || function (thisArg, _arguments, P, generator) {
function adopt(value) { return value instanceof P ? value : new P(function (resolve) { resolve(value); }); }
return new (P || (P = Promise))(function (resolve, reject) {
function fulfilled(value) { try { step(generator.next(value)); } catch (e) { reject(e); } }
function rejected(value) { try { step(generator["throw"](value)); } catch (e) { reject(e); } }
function step(result) { result.done ? resolve(result.value) : adopt(result.value).then(fulfilled, rejected); }
step((generator = generator.apply(thisArg, _arguments || [])).next());
});
};
import path from 'node:path';
import * as utils from './utils.js';
import * as resources from './resources.js';
import { env } from './env.js';
export class GetTaskStatusTool {
constructor() {
this.homeDir = env.POLLO_AI_HOME_DIR;
}
getName() {
return 'getTaskStatus';
}
getDescription() {
return 'Get the status of a generation task.';
}
getInputSchema() {
return {
type: 'object',
properties: {
taskId: {
type: 'string',
description: 'The task id return by generation task.',
},
},
required: ['taskId'],
};
}
execute(request) {
return __awaiter(this, void 0, void 0, function* () {
const { taskId } = request.params.arguments;
const generations = (yield resources.getTaskStatus(taskId)).generations;
if (generations.length === 0) {
return {
content: [{ type: 'text', text: `The task[${taskId}] is not found.` }],
};
}
if (generations.filter((g) => g.status !== 'succeed').length > 0) {
return {
content: [
{ type: 'text', text: `The task[${taskId}] is processing...` },
],
};
}
if (this.homeDir) {
const location = path.join(this.homeDir, taskId);
yield Promise.all(generations.map((g) => utils.downloadFile(location, g.url)));
return {
content: [
{
type: 'text',
text: `The task[${taskId}] generation has been completed. I have downloaded the video to ${location}`,
},
],
};
}
else {
let text = `The task[${taskId}] generation has been completed. You can download the video through the link below, or if you set POLLO_AI_HOME_DIR, I can help you download the video directly to the corresponding directory.\n\t${generations.map((g) => g.url).join('\n\t')}`;
return {
content: [{ type: 'text', text }],
};
}
});
}
}
export class Img2VideoTool {
constructor(modelBrand, modelAlias, schema) {
this.modelBrand = modelBrand;
this.modelAlias = modelAlias;
this.schema = schema;
}
getName() {
return `img2video_${this.modelAlias}`;
}
getDescription() {
return `
Generate a video from a image - Powered by ${this.modelBrand} ${this.modelAlias}.
COST WARNING: This tool makes an API call to Pollo.ai which may incur costs. Only use when explicitly requested by the user.`;
}
getInputSchema() {
return {
type: 'object',
properties: Object.assign({ imagePath: {
type: 'string',
description: 'The local path of the image file to use. If you want to use a remote image, please use the image property directly.',
}, imageTailPath: {
type: 'string',
description: 'The local path of the image tail file to use. If you want to use a remote image, please use imageTail property directly.',
} }, this.schema.properties),
required: [
'imagePath',
...this.schema.required,
].filter((key) => key !== 'image'),
};
}
execute(request) {
return __awaiter(this, void 0, void 0, function* () {
const args = request.params.arguments;
if (!args.image && !args.imagePath) {
return {
content: [
{
type: 'text',
text: 'Please provide an image file path or image url',
},
],
};
}
// upload local images
const uploading = [];
if (args.imagePath) {
uploading.push(utils.uploadImage(args.imagePath));
}
if (args.imageTailPath) {
uploading.push(utils.uploadImage(args.imageTailPath));
}
if (uploading.length) {
const results = yield Promise.all(uploading);
if (results.length === 2) {
args.image = results[0];
args.imageTail = results[1];
}
else if (args.imagePath) {
args.image = results[0];
}
else if (args.imageTailPath) {
args.imageTail = results[0];
}
}
args.imagePath = args.imageTailPath = undefined;
const response = yield resources.generateVideo(this.modelBrand, this.modelAlias, args);
return {
content: [
{
type: 'text',
text: `The generation task status is ${response.status}, and task id is ${response.taskId}.`,
},
],
};
});
}
}
export class Text2VideoTool {
constructor(modelBrand, modelAlias, schema) {
this.modelBrand = modelBrand;
this.modelAlias = modelAlias;
this.schema = schema;
}
getName() {
return `text2video_${this.modelAlias}`;
}
getDescription() {
return `
Generate a video from a text prompt - Powered by ${this.modelBrand} ${this.modelAlias}.
COST WARNING: This tool makes an API call to Pollo.ai which may incur costs. Only use when explicitly requested by the user.`;
}
getInputSchema() {
return this.schema;
}
execute(request) {
return __awaiter(this, void 0, void 0, function* () {
const response = yield resources.generateVideo(this.modelBrand, this.modelAlias, request.params.arguments);
return {
content: [
{
type: 'text',
text: `The generation task status is ${response.status}, and task id is ${response.taskId}.`,
},
],
};
});
}
}
//# sourceMappingURL=tools.js.map