server.ts•5.74 kB
import { Server } from '@modelcontextprotocol/sdk/server/index.js';
import {
CallToolRequestSchema,
ErrorCode,
ListToolsRequestSchema,
McpError,
} from '@modelcontextprotocol/sdk/types.js';
import { HuggingFaceClient } from './client.js';
import { HuggingFaceConfig } from './types.js';
import {
listModelsToolDefinition,
getModelInfoToolDefinition,
getModelTagsToolDefinition,
listDatasetsToolDefinition,
getDatasetInfoToolDefinition,
getDatasetParquetToolDefinition,
getCroissantToolDefinition,
getDatasetTagsToolDefinition,
handleListModels,
handleGetModelInfo,
handleGetModelTags,
handleListDatasets,
handleGetDatasetInfo,
handleGetDatasetParquet,
handleGetCroissant,
handleGetDatasetTags
} from './tools/index.js';
/**
* Main server class for Hugging Face Hub MCP integration
*/
export class HuggingFaceServer {
private client: HuggingFaceClient;
private server: Server;
constructor(config: HuggingFaceConfig) {
this.client = new HuggingFaceClient(config);
this.server = new Server(
{
name: 'huggingface-hub',
version: '0.1.0',
},
{
capabilities: {
tools: {},
},
}
);
this.setupHandlers();
this.setupErrorHandling();
}
private setupHandlers(): void {
this.server.setRequestHandler(ListToolsRequestSchema, async () => ({
tools: [
listModelsToolDefinition,
getModelInfoToolDefinition,
getModelTagsToolDefinition,
listDatasetsToolDefinition,
getDatasetInfoToolDefinition,
getDatasetParquetToolDefinition,
getCroissantToolDefinition,
getDatasetTagsToolDefinition
],
}));
this.server.setRequestHandler(CallToolRequestSchema, async (request) => {
const { name, arguments: args } = request.params;
switch (name) {
case 'hf_list_models':
return handleListModels(this.client, args);
case 'hf_get_model_info':
return handleGetModelInfo(this.client, args);
case 'hf_get_model_tags':
return handleGetModelTags(this.client, args);
case 'hf_list_datasets':
return handleListDatasets(this.client, args);
case 'hf_get_dataset_info':
return handleGetDatasetInfo(this.client, args);
case 'hf_get_dataset_parquet':
return handleGetDatasetParquet(this.client, args);
case 'hf_get_croissant':
return handleGetCroissant(this.client, args);
case 'hf_get_dataset_tags':
return handleGetDatasetTags(this.client, args);
default:
throw new McpError(
ErrorCode.MethodNotFound,
`Unknown tool: ${name}`
);
}
});
}
private setupErrorHandling(): void {
this.server.onerror = (error) => console.error('[MCP Error]', error);
process.on('SIGINT', async () => {
await this.server.close();
process.exit(0);
});
}
getServer(): Server {
return this.server;
}
}
/**
* Factory function for creating standalone server instances
*/
export function createStandaloneServer(config: HuggingFaceConfig): Server {
const server = new Server(
{
name: "huggingface-hub-discovery",
version: "0.1.0",
},
{
capabilities: {
tools: {},
},
},
);
const client = new HuggingFaceClient(config);
server.setRequestHandler(ListToolsRequestSchema, async () => ({
tools: [
listModelsToolDefinition,
getModelInfoToolDefinition,
getModelTagsToolDefinition,
listDatasetsToolDefinition,
getDatasetInfoToolDefinition,
getDatasetParquetToolDefinition,
getCroissantToolDefinition,
getDatasetTagsToolDefinition
],
}));
server.setRequestHandler(CallToolRequestSchema, async (request) => {
const { name, arguments: args } = request.params;
switch (name) {
case 'hf_list_models':
return handleListModels(client, args);
case 'hf_get_model_info':
return handleGetModelInfo(client, args);
case 'hf_get_model_tags':
return handleGetModelTags(client, args);
case 'hf_list_datasets':
return handleListDatasets(client, args);
case 'hf_get_dataset_info':
return handleGetDatasetInfo(client, args);
case 'hf_get_dataset_parquet':
return handleGetDatasetParquet(client, args);
case 'hf_get_croissant':
return handleGetCroissant(client, args);
case 'hf_get_dataset_tags':
return handleGetDatasetTags(client, args);
default:
throw new McpError(
ErrorCode.MethodNotFound,
`Unknown tool: ${name}`
);
}
});
return server;
}