text2sql-mcp-tool
Version:
MCP server for text2sql
287 lines (286 loc) • 10.3 kB
JavaScript
#!/usr/bin/env node
import { Server } from "@modelcontextprotocol/sdk/server/index.js";
import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
import { CallToolRequestSchema, ListToolsRequestSchema, ErrorCode, McpError, } from "@modelcontextprotocol/sdk/types.js";
import { z } from "zod";
import { MastraClient } from "@mastra/client-js";
import axios from "axios";
import fs from "fs/promises";
const QUICKCHART_BASE_URL = "https://quickchart.io/chart";
// Chart types and interfaces
const VALID_CHART_TYPES = [
"bar",
"line",
"pie",
"doughnut",
"radar",
"polarArea",
"scatter",
"bubble",
"radialGauge",
"speedometer",
];
// Create server instance
const server = new Server({
name: "hubble-tool",
version: "1.0.1",
}, {
capabilities: {
tools: {},
},
});
// List available tools
server.setRequestHandler(ListToolsRequestSchema, async () => {
return {
tools: [
{
name: "search-hubble",
description: "Get pumpfun data from Hubble",
inputSchema: {
type: "object",
properties: {
query: {
type: "string",
description: "Search query",
},
},
required: ["query"],
},
},
{
name: "generate_chart",
description: "Generate a chart using QuickChart",
inputSchema: {
type: "object",
properties: {
type: {
type: "string",
description: "Chart type (bar, line, pie, doughnut, radar, polarArea, scatter, bubble, radialGauge, speedometer)",
},
labels: {
type: "array",
items: { type: "string" },
description: "Labels for data points",
},
datasets: {
type: "array",
items: {
type: "object",
properties: {
label: { type: "string" },
data: { type: "array" },
backgroundColor: {
oneOf: [
{ type: "string" },
{ type: "array", items: { type: "string" } },
],
},
borderColor: {
oneOf: [
{ type: "string" },
{ type: "array", items: { type: "string" } },
],
},
additionalConfig: { type: "object" },
},
required: ["data"],
},
},
title: { type: "string" },
options: { type: "object" },
},
required: ["type", "datasets"],
},
},
{
name: "download_chart",
description: "Download a chart image to a local file",
inputSchema: {
type: "object",
properties: {
config: {
type: "object",
description: "Chart configuration object",
},
outputPath: {
type: "string",
description: "Path where the chart image should be saved",
},
},
required: ["config", "outputPath"],
},
},
],
};
});
// Define the schema for search-hubble arguments
const searchHubbleArgsSchema = z.object({
query: z.string(),
});
// Chart utility functions
function validateChartType(type) {
if (!VALID_CHART_TYPES.includes(type)) {
throw new McpError(ErrorCode.InvalidParams, `Invalid chart type. Must be one of: ${VALID_CHART_TYPES.join(", ")}`);
}
}
function generateChartConfig(args) {
const { type, labels, datasets, title, options = {} } = args;
validateChartType(type);
const config = {
type,
data: {
labels: labels || [],
datasets: datasets.map((dataset) => ({
label: dataset.label || "",
data: dataset.data,
backgroundColor: dataset.backgroundColor,
borderColor: dataset.borderColor,
...dataset.additionalConfig,
})),
},
options: {
...options,
...(title && {
title: {
display: true,
text: title,
},
}),
},
};
// Special handling for specific chart types
switch (type) {
case "radialGauge":
case "speedometer":
if (!datasets?.[0]?.data?.[0]) {
throw new McpError(ErrorCode.InvalidParams, `${type} requires a single numeric value`);
}
config.options = {
...config.options,
plugins: {
datalabels: {
display: true,
formatter: (value) => value,
},
},
};
break;
case "scatter":
case "bubble":
datasets.forEach((dataset) => {
if (!Array.isArray(dataset.data[0])) {
throw new McpError(ErrorCode.InvalidParams, `${type} requires data points in [x, y${type === "bubble" ? ", r" : ""}] format`);
}
});
break;
}
return config;
}
async function generateChartUrl(config) {
const encodedConfig = encodeURIComponent(JSON.stringify(config));
return `${QUICKCHART_BASE_URL}?c=${encodedConfig}`;
}
// Handle tool execution
server.setRequestHandler(CallToolRequestSchema, async (request) => {
const { name, arguments: args } = request.params;
try {
if (name === "search-hubble") {
// Validate and parse the arguments with Zod
if (!args) {
throw new Error("Arguments are required for search-hubble");
}
const validatedArgs = searchHubbleArgsSchema.parse(args);
const { query } = validatedArgs;
const client = new MastraClient({
baseUrl: "http://ai-api.hubble-rpc.xyz", // Default Mastra server port
});
const workflow = client.getWorkflow("hubbleWorkflow");
// Start a new workflow run
const result = await workflow.execute({
message: query,
});
console.log(JSON.stringify(result));
return {
content: [
{
type: "text",
text: JSON.stringify(result),
},
],
};
}
else if (name === "generate_chart") {
try {
if (!args) {
throw new McpError(ErrorCode.InvalidParams, "Arguments are required for generate_chart");
}
const config = generateChartConfig(args);
const url = await generateChartUrl(config);
return {
content: [
{
type: "text",
text: url,
},
],
};
}
catch (error) {
if (error instanceof McpError) {
throw error;
}
throw new McpError(ErrorCode.InternalError, `Failed to generate chart: ${error?.message || "Unknown error"}`);
}
}
else if (name === "download_chart") {
try {
if (!args) {
throw new McpError(ErrorCode.InvalidParams, "Arguments are required for download_chart");
}
const { config, outputPath } = args;
const chartConfig = generateChartConfig(config);
const url = await generateChartUrl(chartConfig);
const response = await axios.get(url, { responseType: "arraybuffer" });
await fs.writeFile(outputPath, response.data);
return {
content: [
{
type: "text",
text: `Chart saved to ${outputPath}`,
},
],
};
}
catch (error) {
if (error instanceof McpError) {
throw error;
}
throw new McpError(ErrorCode.InternalError, `Failed to download chart: ${error?.message || "Unknown error"}`);
}
}
else {
throw new Error(`Unknown tool: ${name}`);
}
}
catch (error) {
if (error instanceof z.ZodError) {
throw new Error(`Invalid arguments: ${error.errors
.map((e) => `${e.path.join(".")}: ${e.message}`)
.join(", ")}`);
}
if (error instanceof McpError) {
throw error;
}
throw error;
}
});
// Start the server
async function main() {
const transport = new StdioServerTransport();
await server.connect(transport);
console.error("Hubble MCP Server running on stdio");
}
main().catch((error) => {
console.error("Fatal error in main():", error);
process.exit(1);
});