feat: ConvertToolsToRecords

This commit is contained in:
lilong.129
2025-04-03 15:17:54 +08:00
parent f49ee37055
commit 1081b72847
4 changed files with 380 additions and 2 deletions

2
go.mod
View File

@@ -7,6 +7,7 @@ toolchain go1.23.7
require (
github.com/Masterminds/semver v1.5.0
github.com/andybalholm/brotli v1.0.4
github.com/bytedance/sonic v1.13.2
github.com/cloudwego/eino v0.3.16
github.com/cloudwego/eino-ext/components/model/openai v0.0.0-20250314110024-9e89ba18146c
github.com/cloudwego/eino-ext/components/tool/mcp v0.0.0-20250328102648-b47e7f1587fa
@@ -38,7 +39,6 @@ require (
)
require (
github.com/bytedance/sonic v1.13.2 // indirect
github.com/bytedance/sonic/loader v0.2.4 // indirect
github.com/cenkalti/backoff v2.2.1+incompatible // indirect
github.com/cloudwego/base64x v0.1.5 // indirect

View File

@@ -3,8 +3,11 @@ package mcp
import (
"context"
"fmt"
"strings"
"sync"
"time"
"github.com/bytedance/sonic"
mcpp "github.com/cloudwego/eino-ext/components/tool/mcp"
"github.com/cloudwego/eino/components/tool"
"github.com/httprunner/httprunner/v5/internal/version"
@@ -262,3 +265,149 @@ func (h *MCPHub) CloseServers() error {
return nil
}
// MCPToolRecord represents a single tool record in the database
// Each record contains detailed information about a tool and its server
type MCPToolRecord struct {
ToolID string `json:"tool_id"` // Unique identifier for the tool record
ServerName string `json:"mcp_server"` // Name of the MCP server
ToolName string `json:"tool_name"` // Name of the tool
Description string `json:"description"` // Tool description
Parameters string `json:"parameters"` // Tool input parameters in JSON format
Returns string `json:"returns"` // Tool return value format in JSON format
CreatedAt time.Time `json:"created_at"` // Record creation time
LastUpdatedAt time.Time `json:"last_updated_at"` // Record last update time
}
// DocStringInfo contains the parsed information from a Python docstring
type DocStringInfo struct {
Description string
Parameters map[string]string
Returns map[string]string
}
// extractDocStringInfo extracts information from a Python docstring
// Example input:
// """Get weather alerts for a US state.
//
// Args:
// state: Two-letter US state code (e.g. CA, NY)
//
// Returns:
// alerts: List of active weather alerts for the specified state
// error: Error message if the request fails
// """
func extractDocStringInfo(docstring string) DocStringInfo {
info := DocStringInfo{
Parameters: make(map[string]string),
Returns: make(map[string]string),
}
// Find the Args and Returns sections
argsIndex := strings.Index(docstring, "Args:")
returnsIndex := strings.Index(docstring, "Returns:")
// Extract description (everything before Args)
if argsIndex != -1 {
info.Description = strings.TrimSpace(docstring[:argsIndex])
} else if returnsIndex != -1 {
info.Description = strings.TrimSpace(docstring[:returnsIndex])
} else {
info.Description = strings.TrimSpace(docstring)
return info
}
// Helper function to extract key-value pairs from a section
extractSection := func(content string) map[string]string {
result := make(map[string]string)
lines := strings.Split(content, "\n")
for _, line := range lines {
line = strings.TrimSpace(line)
if line == "" {
continue
}
parts := strings.SplitN(line, ":", 2)
if len(parts) != 2 {
continue
}
key := strings.TrimSpace(parts[0])
value := strings.TrimSpace(parts[1])
if key != "" && value != "" {
result[key] = value
}
}
return result
}
// Extract Args section
if argsIndex != -1 {
endIndex := returnsIndex
if endIndex == -1 {
endIndex = len(docstring)
}
argsContent := docstring[argsIndex+len("Args:") : endIndex]
info.Parameters = extractSection(argsContent)
}
// Extract Returns section
if returnsIndex != -1 {
returnsContent := docstring[returnsIndex+len("Returns:"):]
info.Returns = extractSection(returnsContent)
}
return info
}
// ConvertToolsToRecords converts map[string]MCPTools to a list of database records
func ConvertToolsToRecords(toolsMap map[string]MCPTools) []MCPToolRecord {
var records []MCPToolRecord
now := time.Now()
for serverName, mcpTools := range toolsMap {
if mcpTools.Err != nil {
log.Error().Str("server", serverName).Err(mcpTools.Err).Msg("skip tools conversion due to error")
continue
}
for _, tool := range mcpTools.Tools {
// Generate unique ID by combining server name and tool name
id := fmt.Sprintf("%s_%s", serverName, tool.Name)
// Extract docstring information
info := extractDocStringInfo(tool.Description)
// Convert parameters and returns to JSON
paramsJSON, err := sonic.MarshalString(info.Parameters)
if err != nil {
log.Warn().Interface("params", info.Parameters).Err(err).Msg("failed to marshal parameters to JSON")
paramsJSON = "{}"
}
returnsJSON, err := sonic.MarshalString(info.Returns)
if err != nil {
log.Warn().Interface("returns", info.Returns).Err(err).Msg("failed to marshal returns to JSON")
returnsJSON = "{}"
}
record := MCPToolRecord{
ToolID: id,
ServerName: serverName,
ToolName: tool.Name,
Description: info.Description,
Parameters: paramsJSON,
Returns: returnsJSON,
CreatedAt: now,
LastUpdatedAt: now,
}
records = append(records, record)
}
}
return records
}

View File

@@ -2,9 +2,13 @@ package mcp
import (
"context"
"encoding/json"
"os"
"testing"
"time"
"github.com/cloudwego/eino/components/tool"
"github.com/mark3labs/mcp-go/mcp"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
@@ -54,3 +58,228 @@ func TestCallEinoTool(t *testing.T) {
require.NoError(t, err)
t.Logf("Result: %v", result)
}
func TestConvertToolsToRecordsFromFile(t *testing.T) {
hub, err := NewMCPHub("./test.mcp.json")
require.NoError(t, err)
ctx := context.Background()
err = hub.InitServers(ctx)
require.NoError(t, err)
tools := hub.GetTools(ctx)
require.NoError(t, err)
records := ConvertToolsToRecords(tools)
// Convert records to JSON
recordsJSON, err := json.Marshal(records)
require.NoError(t, err)
// Write JSON to file
err = os.WriteFile("./tools_records.json", recordsJSON, 0o644)
require.NoError(t, err)
t.Logf("Tools records written to ./tools_records.json")
}
func TestExtractDocStringInfo(t *testing.T) {
tests := []struct {
name string
docstring string
want DocStringInfo
}{
{
name: "complete docstring with args and returns",
docstring: `Get weather alerts for a US state.
Args:
state: Two-letter US state code (e.g. CA, NY)
Returns:
alerts: List of active weather alerts for the specified state
error: Error message if the request fails
`,
want: DocStringInfo{
Description: "Get weather alerts for a US state.",
Parameters: map[string]string{
"state": "Two-letter US state code (e.g. CA, NY)",
},
Returns: map[string]string{
"alerts": "List of active weather alerts for the specified state",
"error": "Error message if the request fails",
},
},
},
{
name: "docstring with only args",
docstring: `Do screen swipe action.
Args:
direction: swipe direction (up, down)
`,
want: DocStringInfo{
Description: "Do screen swipe action.",
Parameters: map[string]string{
"direction": "swipe direction (up, down)",
},
Returns: map[string]string{},
},
},
{
name: "docstring with only description",
docstring: "Simple tool with no parameters.",
want: DocStringInfo{
Description: "Simple tool with no parameters.",
Parameters: map[string]string{},
Returns: map[string]string{},
},
},
{
name: "docstring with multiple parameters",
docstring: `Perform complex operation.
Args:
param1: first parameter description
param2: second parameter description
param3: third parameter description
Returns:
result: operation result
`,
want: DocStringInfo{
Description: "Perform complex operation.",
Parameters: map[string]string{
"param1": "first parameter description",
"param2": "second parameter description",
"param3": "third parameter description",
},
Returns: map[string]string{
"result": "operation result",
},
},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := extractDocStringInfo(tt.docstring)
assert.Equal(t, tt.want.Description, got.Description)
assert.Equal(t, tt.want.Parameters, got.Parameters)
assert.Equal(t, tt.want.Returns, got.Returns)
})
}
}
func TestConvertToolsToRecords(t *testing.T) {
tests := []struct {
name string
toolsMap map[string]MCPTools
want []MCPToolRecord
}{
{
name: "convert weather tool",
toolsMap: map[string]MCPTools{
"weather": {
Name: "weather",
Tools: []mcp.Tool{
{
Name: "get_alerts",
Description: `Get weather alerts for a US state.
Args:
state: Two-letter US state code (e.g. CA, NY)
Returns:
alerts: List of active weather alerts for the specified state
error: Error message if the request fails
`,
},
},
},
},
want: []MCPToolRecord{
{
ToolID: "weather_get_alerts",
ServerName: "weather",
ToolName: "get_alerts",
Description: "Get weather alerts for a US state.",
Parameters: `{"state":"Two-letter US state code (e.g. CA, NY)"}`,
Returns: `{"alerts":"List of active weather alerts for the specified state","error":"Error message if the request fails"}`,
},
},
},
{
name: "convert multiple tools",
toolsMap: map[string]MCPTools{
"ui": {
Name: "ui",
Tools: []mcp.Tool{
{
Name: "swipe",
Description: `Do screen swipe action.
Args:
direction: swipe direction (up, down)
`,
},
{
Name: "tap",
Description: "Tap on screen at specified position.",
},
},
},
},
want: []MCPToolRecord{
{
ToolID: "ui_swipe",
ServerName: "ui",
ToolName: "swipe",
Description: "Do screen swipe action.",
Parameters: `{"direction":"swipe direction (up, down)"}`,
Returns: "{}",
},
{
ToolID: "ui_tap",
ServerName: "ui",
ToolName: "tap",
Description: "Tap on screen at specified position.",
Parameters: "{}",
Returns: "{}",
},
},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := ConvertToolsToRecords(tt.toolsMap)
// Compare each record
require.Equal(t, len(tt.want), len(got))
for i := range tt.want {
assert.Equal(t, tt.want[i].ToolID, got[i].ToolID)
assert.Equal(t, tt.want[i].ServerName, got[i].ServerName)
assert.Equal(t, tt.want[i].ToolName, got[i].ToolName)
assert.Equal(t, tt.want[i].Description, got[i].Description)
// Compare JSON content (ignoring whitespace differences)
var wantParams, gotParams, wantReturns, gotReturns map[string]string
require.NoError(t, json.Unmarshal([]byte(tt.want[i].Parameters), &wantParams))
require.NoError(t, json.Unmarshal([]byte(got[i].Parameters), &gotParams))
require.NoError(t, json.Unmarshal([]byte(tt.want[i].Returns), &wantReturns))
require.NoError(t, json.Unmarshal([]byte(got[i].Returns), &gotReturns))
assert.Equal(t, wantParams, gotParams)
assert.Equal(t, wantReturns, gotReturns)
// Verify timestamps are recent (within last 5 seconds)
now := time.Now()
assert.True(t, now.Sub(got[i].CreatedAt) < 5*time.Second, "CreatedAt should be recent")
assert.True(t, now.Sub(got[i].LastUpdatedAt) < 5*time.Second, "LastUpdatedAt should be recent")
// CreatedAt and LastUpdatedAt should be the same
assert.Equal(t, got[i].CreatedAt, got[i].LastUpdatedAt)
}
})
}
}

View File

@@ -1 +1 @@
v5.0.0-beta-2504012233
v5.0.0-beta-2504031520