titor
d3e04c5d1a
- 支持三种图表: 折线图、柱状图、散点图 - MCP 协议支持 (stdio + HTTP) - 完整的单元测试和集成测试 - Docker 支持 - Makefile 构建脚本
338 lines
8.6 KiB
Go
338 lines
8.6 KiB
Go
// Package tools ttychart-mcp 工具实现包
|
||
//
|
||
// 该包提供所有 MCP 工具的具体实现，包括:
|
||
// - plot_line: 折线图
|
||
// - plot_bar: 柱状图
|
||
// - plot_scatter: 散点图
|
||
//
|
||
// 作者: titor
|
||
// 创建日期: 2026-04-15
|
||
package tools
|
||
|
||
import (
|
||
"context"
|
||
"encoding/json"
|
||
"fmt"
|
||
"strconv"
|
||
"strings"
|
||
|
||
"hub.gaomia.site/titor/ttychart-mcp/tools/charts"
|
||
|
||
"charm.land/lipgloss/v2"
|
||
"github.com/mark3labs/mcp-go/mcp"
|
||
"github.com/mark3labs/mcp-go/server"
|
||
)
|
||
|
||
// NewServer 创建新的 MCP 服务器
|
||
func NewServer() *server.MCPServer {
|
||
srv := server.NewMCPServer(
|
||
"ttychart",
|
||
"0.1.0",
|
||
server.WithToolCapabilities(true),
|
||
)
|
||
|
||
// 折线图工具
|
||
plotLineTool := mcp.NewTool(
|
||
"plot_line",
|
||
mcp.WithDescription("绘制折线图，用于展示数据随时间或类别的变化趋势"),
|
||
mcp.WithNumber("data",
|
||
mcp.Required(),
|
||
mcp.Description("Y轴数据点数组"),
|
||
),
|
||
mcp.WithString("title",
|
||
mcp.Description("图表标题"),
|
||
),
|
||
mcp.WithNumber("width",
|
||
mcp.Description("图表宽度（字符数），默认自动计算"),
|
||
),
|
||
mcp.WithNumber("height",
|
||
mcp.Description("图表高度（行数），默认自动计算"),
|
||
),
|
||
mcp.WithString("labels",
|
||
mcp.Description("X轴标签，多个标签用逗号分隔"),
|
||
),
|
||
mcp.WithString("color",
|
||
mcp.Description("线条颜色，可选值: red, green, blue, yellow, cyan, magenta"),
|
||
),
|
||
)
|
||
|
||
// 柱状图工具
|
||
plotBarTool := mcp.NewTool(
|
||
"plot_bar",
|
||
mcp.WithDescription("绘制柱状图，用于展示分类数据的比较"),
|
||
mcp.WithString("data",
|
||
mcp.Required(),
|
||
mcp.Description("数据，格式: 标签1:数值1,标签2:数值2"),
|
||
),
|
||
mcp.WithString("title",
|
||
mcp.Description("图表标题"),
|
||
),
|
||
mcp.WithNumber("width",
|
||
mcp.Description("图表宽度"),
|
||
),
|
||
mcp.WithNumber("height",
|
||
mcp.Description("图表高度"),
|
||
),
|
||
mcp.WithBoolean("horizontal",
|
||
mcp.Description("是否为水平柱状图，默认为垂直"),
|
||
),
|
||
mcp.WithString("color",
|
||
mcp.Description("柱体颜色"),
|
||
),
|
||
)
|
||
|
||
// 散点图工具
|
||
plotScatterTool := mcp.NewTool(
|
||
"plot_scatter",
|
||
mcp.WithDescription("绘制散点图，用于展示数据的相关性"),
|
||
mcp.WithString("data",
|
||
mcp.Required(),
|
||
mcp.Description("数据点，格式: x1,y1 x2,y2 ..."),
|
||
),
|
||
mcp.WithString("title",
|
||
mcp.Description("图表标题"),
|
||
),
|
||
mcp.WithNumber("width",
|
||
mcp.Description("图表宽度"),
|
||
),
|
||
mcp.WithNumber("height",
|
||
mcp.Description("图表高度"),
|
||
),
|
||
mcp.WithString("x_label",
|
||
mcp.Description("X轴标签"),
|
||
),
|
||
mcp.WithString("y_label",
|
||
mcp.Description("Y轴标签"),
|
||
),
|
||
mcp.WithString("color",
|
||
mcp.Description("点颜色"),
|
||
),
|
||
)
|
||
|
||
srv.AddTool(plotLineTool, handlePlotLineTool)
|
||
srv.AddTool(plotBarTool, handlePlotBarTool)
|
||
srv.AddTool(plotScatterTool, handlePlotScatterTool)
|
||
|
||
return srv
|
||
}
|
||
|
||
// handlePlotLineTool 处理折线图工具调用
|
||
func handlePlotLineTool(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
|
||
// 将请求参数转换为 JSON 再解析
|
||
jsonBytes, err := json.Marshal(request.Params)
|
||
if err != nil {
|
||
return mcp.NewToolResultError("参数解析失败"), err
|
||
}
|
||
|
||
var args struct {
|
||
Data interface{} `json:"data"`
|
||
Title string `json:"title"`
|
||
Width int `json:"width"`
|
||
Height int `json:"height"`
|
||
Labels string `json:"labels"`
|
||
Color string `json:"color"`
|
||
}
|
||
|
||
if err := json.Unmarshal(jsonBytes, &args); err != nil {
|
||
return mcp.NewToolResultError("参数解析失败"), err
|
||
}
|
||
|
||
// 处理数据
|
||
var data []float64
|
||
switch v := args.Data.(type) {
|
||
case float64:
|
||
data = []float64{v}
|
||
case []interface{}:
|
||
data = make([]float64, 0, len(v))
|
||
for _, item := range v {
|
||
if f, ok := toFloat64(item); ok {
|
||
data = append(data, f)
|
||
}
|
||
}
|
||
}
|
||
|
||
if len(data) == 0 {
|
||
return mcp.NewToolResultError("错误: data 参数无效"), nil
|
||
}
|
||
|
||
// 解析标签
|
||
var labels []string
|
||
if args.Labels != "" {
|
||
labels = strings.Split(args.Labels, ",")
|
||
}
|
||
|
||
// 获取颜色
|
||
color := getColorStyle(args.Color)
|
||
|
||
// 生成图表
|
||
graph := charts.NewLineChart(args.Width, args.Height)
|
||
graph.SetData(data)
|
||
graph.SetTitle(args.Title)
|
||
graph.SetLabels(labels)
|
||
graph.SetColor(color)
|
||
|
||
return mcp.NewToolResultText(graph.Render()), nil
|
||
}
|
||
|
||
// handlePlotBarTool 处理柱状图工具调用
|
||
func handlePlotBarTool(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
|
||
jsonBytes, err := json.Marshal(request.Params)
|
||
if err != nil {
|
||
return mcp.NewToolResultError("参数解析失败"), err
|
||
}
|
||
|
||
var args struct {
|
||
Data interface{} `json:"data"`
|
||
Title string `json:"title"`
|
||
Width int `json:"width"`
|
||
Height int `json:"height"`
|
||
Horizontal bool `json:"horizontal"`
|
||
Color string `json:"color"`
|
||
}
|
||
|
||
if err := json.Unmarshal(jsonBytes, &args); err != nil {
|
||
return mcp.NewToolResultError("参数解析失败"), err
|
||
}
|
||
|
||
dataStr, ok := args.Data.(string)
|
||
if !ok {
|
||
return mcp.NewToolResultError("错误: data 参数格式无效"), nil
|
||
}
|
||
|
||
// 解析数据: "标签1:数值1,标签2:数值2"
|
||
data := parseBarData(dataStr)
|
||
if len(data) == 0 {
|
||
return mcp.NewToolResultError("错误: data 参数格式无效"), nil
|
||
}
|
||
|
||
color := getColorStyle(args.Color)
|
||
|
||
graph := charts.NewBarChart(args.Width, args.Height, args.Horizontal)
|
||
graph.SetData(data)
|
||
graph.SetTitle(args.Title)
|
||
graph.SetColor(color)
|
||
|
||
return mcp.NewToolResultText(graph.Render()), nil
|
||
}
|
||
|
||
// handlePlotScatterTool 处理散点图工具调用
|
||
func handlePlotScatterTool(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
|
||
jsonBytes, err := json.Marshal(request.Params)
|
||
if err != nil {
|
||
return mcp.NewToolResultError("参数解析失败"), err
|
||
}
|
||
|
||
var args struct {
|
||
Data interface{} `json:"data"`
|
||
Title string `json:"title"`
|
||
Width int `json:"width"`
|
||
Height int `json:"height"`
|
||
XLabel string `json:"x_label"`
|
||
YLabel string `json:"y_label"`
|
||
Color string `json:"color"`
|
||
}
|
||
|
||
if err := json.Unmarshal(jsonBytes, &args); err != nil {
|
||
return mcp.NewToolResultError("参数解析失败"), err
|
||
}
|
||
|
||
dataStr, ok := args.Data.(string)
|
||
if !ok {
|
||
return mcp.NewToolResultError("错误: data 参数格式无效"), nil
|
||
}
|
||
|
||
// 解析数据: "x1,y1 x2,y2 ..."
|
||
data := parseScatterData(dataStr)
|
||
if len(data) == 0 {
|
||
return mcp.NewToolResultError("错误: data 参数格式无效"), nil
|
||
}
|
||
|
||
color := getColorStyle(args.Color)
|
||
|
||
graph := charts.NewScatterChart(args.Width, args.Height)
|
||
graph.SetData(data)
|
||
graph.SetTitle(args.Title)
|
||
graph.SetXLabel(args.XLabel)
|
||
graph.SetYLabel(args.YLabel)
|
||
graph.SetColor(color)
|
||
|
||
return mcp.NewToolResultText(graph.Render()), nil
|
||
}
|
||
|
||
// parseBarData 解析柱状图数据
|
||
// 格式: "标签1:数值1,标签2:数值2"
|
||
func parseBarData(str string) []charts.BarDataItem {
|
||
var result []charts.BarDataItem
|
||
items := strings.Split(str, ",")
|
||
for _, item := range items {
|
||
parts := strings.SplitN(item, ":", 2)
|
||
if len(parts) == 2 {
|
||
if value, err := strconv.ParseFloat(strings.TrimSpace(parts[1]), 64); err == nil {
|
||
result = append(result, charts.BarDataItem{
|
||
Label: strings.TrimSpace(parts[0]),
|
||
Value: value,
|
||
})
|
||
}
|
||
}
|
||
}
|
||
return result
|
||
}
|
||
|
||
// parseScatterData 解析散点图数据
|
||
// 格式: "x1,y1 x2,y2 ..."
|
||
func parseScatterData(str string) []charts.ScatterDataItem {
|
||
var result []charts.ScatterDataItem
|
||
items := strings.Fields(str)
|
||
for _, item := range items {
|
||
parts := strings.SplitN(item, ",", 2)
|
||
if len(parts) == 2 {
|
||
x, err1 := strconv.ParseFloat(strings.TrimSpace(parts[0]), 64)
|
||
y, err2 := strconv.ParseFloat(strings.TrimSpace(parts[1]), 64)
|
||
if err1 == nil && err2 == nil {
|
||
result = append(result, charts.ScatterDataItem{
|
||
X: x,
|
||
Y: y,
|
||
})
|
||
}
|
||
}
|
||
}
|
||
return result
|
||
}
|
||
|
||
// getColorStyle 获取颜色样式
|
||
func getColorStyle(colorName string) lipgloss.Style {
|
||
switch colorName {
|
||
case "red":
|
||
return lipgloss.NewStyle().Foreground(lipgloss.Color("9"))
|
||
case "green":
|
||
return lipgloss.NewStyle().Foreground(lipgloss.Color("10"))
|
||
case "blue":
|
||
return lipgloss.NewStyle().Foreground(lipgloss.Color("12"))
|
||
case "yellow":
|
||
return lipgloss.NewStyle().Foreground(lipgloss.Color("11"))
|
||
case "cyan":
|
||
return lipgloss.NewStyle().Foreground(lipgloss.Color("14"))
|
||
case "magenta":
|
||
return lipgloss.NewStyle().Foreground(lipgloss.Color("13"))
|
||
default:
|
||
return lipgloss.NewStyle().Foreground(lipgloss.Color("15"))
|
||
}
|
||
}
|
||
|
||
// toFloat64 将 interface{} 转换为 float64
|
||
func toFloat64(v interface{}) (float64, bool) {
|
||
switch n := v.(type) {
|
||
case float64:
|
||
return n, true
|
||
case int:
|
||
return float64(n), true
|
||
case int64:
|
||
return float64(n), true
|
||
default:
|
||
return 0, false
|
||
}
|
||
}
|
||
|
||
// 这里添加 fmt 的使用以避免导入错误
|
||
var _ = fmt.Sprintf
|