huggingface-mcp-server
Version:
MCP Server for HuggingFace inference endpoints with custom LoRA and story generation
504 lines (503 loc) • 24 kB
JavaScript
;
var __createBinding = (this && this.__createBinding) || (Object.create ? (function(o, m, k, k2) {
if (k2 === undefined) k2 = k;
var desc = Object.getOwnPropertyDescriptor(m, k);
if (!desc || ("get" in desc ? !m.__esModule : desc.writable || desc.configurable)) {
desc = { enumerable: true, get: function() { return m[k]; } };
}
Object.defineProperty(o, k2, desc);
}) : (function(o, m, k, k2) {
if (k2 === undefined) k2 = k;
o[k2] = m[k];
}));
var __setModuleDefault = (this && this.__setModuleDefault) || (Object.create ? (function(o, v) {
Object.defineProperty(o, "default", { enumerable: true, value: v });
}) : function(o, v) {
o["default"] = v;
});
var __importStar = (this && this.__importStar) || (function () {
var ownKeys = function(o) {
ownKeys = Object.getOwnPropertyNames || function (o) {
var ar = [];
for (var k in o) if (Object.prototype.hasOwnProperty.call(o, k)) ar[ar.length] = k;
return ar;
};
return ownKeys(o);
};
return function (mod) {
if (mod && mod.__esModule) return mod;
var result = {};
if (mod != null) for (var k = ownKeys(mod), i = 0; i < k.length; i++) if (k[i] !== "default") __createBinding(result, mod, k[i]);
__setModuleDefault(result, mod);
return result;
};
})();
Object.defineProperty(exports, "__esModule", { value: true });
exports.startStdioServer = startStdioServer;
const readline = __importStar(require("readline"));
const handler_1 = require("./handler");
/**
* Starts a server that communicates over standard input/output using JSON-RPC 2.0
* @param apiKey The HuggingFace API key
* @returns A Promise that resolves when the server is shut down
*/
async function startStdioServer(apiKey) {
// --------- LOGGING SETTINGS ---------
// Set to true for verbose logging, false for minimal logging
const VERBOSE_LOGGING = false;
// Custom logging function that respects verbosity setting
const log = (message, isError = false) => {
if (VERBOSE_LOGGING || isError) {
console.error(message);
}
};
// Use stderr for logging in stdio mode
log('Starting HuggingFace MCP Server in stdio mode', true);
log(`API Key is ${apiKey ? 'provided' : 'missing'}`, true);
// Create readline interface for stdio communication
const rl = readline.createInterface({
input: process.stdin,
output: process.stdout,
terminal: false
});
// Define server info and tools
const serverInfo = {
name: 'HuggingFace MCP Server',
version: '1.0.18'
};
// Define tools following the MCP tool definition structure
const tools = [
{
name: 'generate_image', // Name is the unique identifier
description: 'Generate an image based on a text prompt using Flux model, with optional custom LoRA',
inputSchema: {
type: 'object',
properties: {
prompt: {
type: 'string',
description: 'Description of the image to generate'
},
num_inference_steps: {
type: 'number',
description: 'The number of denoising steps of the image. More steps usually lead to higher quality at the cost of slower inference',
default: 25
},
height: {
type: 'number',
description: 'Height of the image in pixels',
default: 1024
},
width: {
type: 'number',
description: 'Width of the image in pixels',
default: 1024
},
guidance_scale: {
type: 'number',
description: 'Controls how much the image generation follows the text prompt. Higher values make the image stick more closely to the input text',
default: 3.5
},
seed: {
type: 'number',
description: 'A starting point to initiate the generation process, use 0 for a random seed',
default: 0
},
num_images_per_prompt: {
type: 'number',
description: 'Number of images to generate with the settings',
default: 1
},
lora_name: {
type: 'string',
description: 'Name of the custom LoRA model on HuggingFace (e.g., "username/lora-model-name")'
}
},
required: ['prompt']
}
},
{
name: 'generate_story',
description: 'Generate a story based on a prompt',
inputSchema: {
type: 'object',
properties: {
prompt: {
type: 'string',
description: 'Prompt for story generation'
}
},
required: ['prompt']
}
}
];
// Process each line from stdin as a JSON-RPC request
return new Promise((resolve) => {
rl.on('line', async (line) => {
try {
// Parse the incoming JSON-RPC message
const request = JSON.parse(line);
// Log the request for debugging
log(`Received request: ${JSON.stringify({
method: request.method,
id: request.id,
paramsSize: request.params ? Object.keys(request.params).length : 0
})}`, false);
// Validate JSON-RPC 2.0 request
if (request.jsonrpc !== '2.0') {
log('Received non-2.0 JSON-RPC request', true);
console.log(JSON.stringify({
jsonrpc: '2.0',
id: request.id || null,
error: {
code: -32600,
message: 'Invalid Request: jsonrpc must be "2.0"'
}
}));
return;
}
// Handle initialization requests and notifications
if (request.method === 'initialize') {
// Handle initialize request from client
const clientInfo = request.params?.clientInfo || { name: 'Unknown Client', version: '0.0.0' };
const clientProtocolVersion = request.params?.protocolVersion || 'unknown';
log(`Received initialize request from ${clientInfo.name} ${clientInfo.version} with protocol version ${clientProtocolVersion}`, false);
// Create tool capabilities object - each tool should be a key in the object
const toolCapabilities = {};
tools.forEach(tool => {
toolCapabilities[tool.name] = true;
});
// Respond with server info, protocol version, and capabilities
console.log(JSON.stringify({
jsonrpc: '2.0',
id: request.id,
result: {
serverInfo,
protocolVersion: '2024-11-05', // Updated to match the client's expected protocol version
capabilities: {
tools: toolCapabilities,
chat: true,
resources: {} // Changed from boolean to empty object
}
}
}));
}
else if (request.method === 'initialized') {
// Handle initialized notification from client
log('Received initialized notification, initialization complete', false);
// No response needed for notification
}
else if (request.method === 'notifications/initialized') {
// Handle notifications/initialized notification from client
log('Received notifications/initialized notification', false);
// No response needed for notification
}
else if (request.method === 'tools/list' || request.method === 'listTools' || request.method === 'tools') {
// Handle tools listing requests (supporting all variations of endpoint names)
log(`Handling ${request.method} request`, false);
console.log(JSON.stringify({
jsonrpc: '2.0',
id: request.id,
result: {
tools: tools.map(tool => ({
name: tool.name,
description: tool.description,
inputSchema: tool.inputSchema
}))
}
}));
}
else if (request.method === 'resources/list') {
// Handle resources listing request
log('Handling resources/list request', false);
// Define available resources - add your actual resources here
const resources = [
{
uri: 'huggingface://models/recent',
name: 'Recent HuggingFace Models',
description: 'List of recently used HuggingFace models',
mimeType: 'text/plain'
}
];
console.log(JSON.stringify({
jsonrpc: '2.0',
id: request.id,
result: {
resources: resources
}
}));
}
else if (request.method === 'resources/templates/list') {
// Handle resource templates listing request
log('Handling resources/templates/list request', false);
// Define available resource templates - add your actual templates here
const resourceTemplates = [
{
name: 'huggingface-model',
description: 'Retrieve information about a HuggingFace model',
uriTemplate: 'huggingface://models/{model_id}',
parameters: [
{
name: 'model_id',
description: 'The ID of the HuggingFace model',
required: true
}
]
}
];
console.log(JSON.stringify({
jsonrpc: '2.0',
id: request.id,
result: {
templates: resourceTemplates
}
}));
}
else if (request.method === 'tools/call') {
// Handle tools/call request
log(`Handling tools/call request for tool: ${request.params?.name}`, false);
if (!request.params?.name) {
console.log(JSON.stringify({
jsonrpc: '2.0',
id: request.id,
error: {
code: -32602,
message: 'Invalid params: tool name is required'
}
}));
return;
}
// Convert the tools/call request to a tool_call format our handler understands
const toolCall = {
id: `call_${Date.now()}`,
type: 'function',
function: {
name: request.params.name,
arguments: JSON.stringify(request.params.arguments || {})
}
};
try {
const results = await (0, handler_1.handleToolCalls)([toolCall], apiKey);
// Format as MCP response with the right content structure
if (results && results.length > 0) {
const toolResult = results[0];
// Check if the result contains an image (basic check for base64 data)
const isImage = toolResult.content.includes('data:image');
if (isImage) {
// Extract the base64 data from the response
const base64Match = toolResult.content.match(/data:image\/[^;]+;base64,([^"]+)/);
const base64Data = base64Match ? base64Match[1] : '';
const mimeType = 'image/jpeg'; // Assume JPEG for simplicity
console.log(JSON.stringify({
jsonrpc: '2.0',
id: request.id,
result: {
content: [
{
type: 'text',
text: toolResult.content.split('\n')[0] // Just the first line as text
},
{
type: 'image',
data: base64Data,
mimeType: mimeType
}
]
}
}));
}
else {
// For non-image responses, just return text
console.log(JSON.stringify({
jsonrpc: '2.0',
id: request.id,
result: {
content: [{
type: 'text',
text: toolResult.content
}]
}
}));
}
}
else {
console.log(JSON.stringify({
jsonrpc: '2.0',
id: request.id,
error: {
code: -32000,
message: 'No result returned from tool call'
}
}));
}
}
catch (error) {
log(`Error in tools/call: ${error.message}`, true);
console.log(JSON.stringify({
jsonrpc: '2.0',
id: request.id,
error: {
code: -32000,
message: `Error processing tool call: ${error.message}`
}
}));
}
}
else if (request.method === 'chat') {
// Handle chat completion request
// Handle both direct and wrapped formats for chat requests
const params = request.params || {};
const chatRequest = params.data;
const messages = chatRequest?.messages || params.messages;
log(`Chat request received with ${messages?.length || 0} messages`, false);
if (!messages || !Array.isArray(messages)) {
console.log(JSON.stringify({
jsonrpc: '2.0',
id: request.id,
error: {
code: -32602,
message: 'Invalid params: messages is required and must be an array'
}
}));
return;
}
// Check if the last message contains tool calls
const lastMessage = messages[messages.length - 1];
if (lastMessage.role === 'assistant' && lastMessage.tool_calls) {
try {
const results = await (0, handler_1.handleToolCalls)(lastMessage.tool_calls, apiKey);
const response = {
jsonrpc: '2.0',
id: request.id,
result: {
id: 'chatcmpl-456',
object: 'chat.completion',
created: Math.floor(Date.now() / 1000),
model: 'hf-mcp-server',
choices: [{
index: 0,
message: {
role: 'assistant',
content: 'I\'ve processed your request.',
tool_calls: []
},
finish_reason: 'tool_calls'
}],
tool_responses: results
}
};
console.log(JSON.stringify(response));
}
catch (error) {
console.log(JSON.stringify({
jsonrpc: '2.0',
id: request.id,
error: {
code: -32000,
message: `Error processing tool calls: ${error.message}`
}
}));
}
}
else {
// Return a response asking the user to choose a tool
const response = {
jsonrpc: '2.0',
id: request.id,
result: {
id: 'chatcmpl-123',
object: 'chat.completion',
created: Math.floor(Date.now() / 1000),
model: 'hf-mcp-server',
choices: [{
index: 0,
message: {
role: 'assistant',
tool_calls: [
{
id: 'call_abc123',
type: 'function',
function: {
name: 'generate_image',
arguments: JSON.stringify({
prompt: 'I need a description from you for the image you\'d like to generate.'
})
}
}
]
},
finish_reason: 'tool_calls'
}]
}
};
console.log(JSON.stringify(response));
}
}
else if (request.method === 'exit') {
// Handle exit request
console.log(JSON.stringify({
jsonrpc: '2.0',
id: request.id,
result: { success: true }
}));
rl.close();
resolve();
}
else if (request.method === 'shutdown') {
// Handle shutdown request
console.log(JSON.stringify({
jsonrpc: '2.0',
id: request.id,
result: { success: true }
}));
rl.close();
resolve();
}
else {
// Handle unknown method
log(`Unknown method: ${request.method}`, true);
console.log(JSON.stringify({
jsonrpc: '2.0',
id: request.id,
error: {
code: -32601,
message: `Method not found: ${request.method}`,
data: {
supportedMethods: [
'initialize',
'initialized',
'notifications/initialized',
'tools/list',
'listTools',
'tools',
'resources/list',
'resources/templates/list',
'tools/call',
'chat/completions',
'shutdown'
],
serverName: serverInfo.name,
serverVersion: serverInfo.version
}
}
}));
}
}
catch (error) {
// Handle JSON parsing errors or other exceptions
log(`Error processing request: ${error.message}`, true);
console.log(JSON.stringify({
jsonrpc: '2.0',
id: null,
error: {
code: -32700,
message: `Parse error: ${error.message}`
}
}));
}
});
// Handle readline close event
rl.on('close', () => {
resolve();
});
});
}