mirror of
https://github.com/httprunner/httprunner.git
synced 2026-08-28 19:47:14 +08:00
feat: 重构 MCP 工具导出逻辑并完善返回值类型系统
This commit is contained in:
lilong.129
+135
-40
@@ -9,20 +9,30 @@ import (
|
||||
"time"
|
||||
|
||||
"github.com/bytedance/sonic"
|
||||
"github.com/mark3labs/mcp-go/mcp"
|
||||
"github.com/rs/zerolog/log"
|
||||
|
||||
"github.com/httprunner/httprunner/v5/uixt"
|
||||
"github.com/httprunner/httprunner/v5/uixt/option"
|
||||
)
|
||||
|
||||
// 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
|
||||
ToolID string `json:"tool_id"` // Unique identifier for the tool record
|
||||
BizID string `json:"biz_id"` // Business ID of the tool
|
||||
VisibleRange int `json:"visible_range"` // Visible range of the tool, 0: visible to biz, 1: visible to all
|
||||
ToolType string `json:"tool_type"` // Type of the tool
|
||||
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:"return_desc"` // Tool return value format in JSON format
|
||||
TeardownPair string `json:"teardown_pair"` // Teardown pair of the tool
|
||||
Examples string `json:"examples"` // Examples of the tool
|
||||
SupportPatterns string `json:"support_patterns"` // Support pattern of the tool
|
||||
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
|
||||
@@ -109,8 +119,13 @@ func extractDocStringInfo(docstring string) DocStringInfo {
|
||||
return info
|
||||
}
|
||||
|
||||
// ActionToolProvider defines the interface for MCP servers that provide ActionTool implementations
|
||||
type ActionToolProvider interface {
|
||||
GetToolByAction(actionName option.ActionName) uixt.ActionTool
|
||||
}
|
||||
|
||||
// ConvertToolsToRecords converts []MCPTools to a list of database records
|
||||
func ConvertToolsToRecords(tools []MCPTools) []MCPToolRecord {
|
||||
func (host *MCPHost) ConvertToolsToRecords(tools []MCPTools) []MCPToolRecord {
|
||||
var records []MCPToolRecord
|
||||
now := time.Now()
|
||||
|
||||
@@ -121,36 +136,7 @@ func ConvertToolsToRecords(tools []MCPTools) []MCPToolRecord {
|
||||
}
|
||||
|
||||
for _, tool := range mcpTools.Tools {
|
||||
// Generate unique ID by combining server name and tool name
|
||||
id := fmt.Sprintf("%s__%s", mcpTools.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: mcpTools.ServerName,
|
||||
ToolName: tool.Name,
|
||||
Description: info.Description,
|
||||
Parameters: paramsJSON,
|
||||
Returns: returnsJSON,
|
||||
CreatedAt: now,
|
||||
LastUpdatedAt: now,
|
||||
}
|
||||
|
||||
record := host.convertSingleToolToRecord(mcpTools.ServerName, tool, now)
|
||||
records = append(records, record)
|
||||
}
|
||||
}
|
||||
@@ -158,12 +144,121 @@ func ConvertToolsToRecords(tools []MCPTools) []MCPToolRecord {
|
||||
return records
|
||||
}
|
||||
|
||||
// convertSingleToolToRecord converts a single MCP tool to a database record
|
||||
func (host *MCPHost) convertSingleToolToRecord(serverName string, tool mcp.Tool, timestamp time.Time) MCPToolRecord {
|
||||
// Generate unique ID
|
||||
id := fmt.Sprintf("%s__%s", serverName, tool.Name)
|
||||
|
||||
// Extract description from docstring
|
||||
info := extractDocStringInfo(tool.Description)
|
||||
|
||||
// Extract parameters
|
||||
paramsJSON := host.extractParameters(tool, info)
|
||||
|
||||
// Extract returns
|
||||
returnsJSON := host.extractReturns(serverName, tool.Name, info)
|
||||
|
||||
return MCPToolRecord{
|
||||
ToolID: id,
|
||||
VisibleRange: 1,
|
||||
ToolType: "edge",
|
||||
ServerName: serverName,
|
||||
ToolName: tool.Name,
|
||||
Description: info.Description,
|
||||
Parameters: paramsJSON,
|
||||
Returns: returnsJSON,
|
||||
CreatedAt: timestamp,
|
||||
LastUpdatedAt: timestamp,
|
||||
}
|
||||
}
|
||||
|
||||
// extractParameters extracts parameter information from tool schema or docstring
|
||||
func (host *MCPHost) extractParameters(tool mcp.Tool, info DocStringInfo) string {
|
||||
// Priority 1: Extract from InputSchema.Properties
|
||||
if len(tool.InputSchema.Properties) > 0 {
|
||||
return host.extractParametersFromSchema(tool.InputSchema.Properties)
|
||||
}
|
||||
|
||||
// Priority 2: Extract from docstring
|
||||
if len(info.Parameters) > 0 {
|
||||
return host.marshalToJSON(info.Parameters, "docstring parameters")
|
||||
}
|
||||
|
||||
return "{}"
|
||||
}
|
||||
|
||||
// extractParametersFromSchema extracts parameters from MCP tool input schema
|
||||
func (host *MCPHost) extractParametersFromSchema(properties map[string]interface{}) string {
|
||||
schemaParams := make(map[string]string)
|
||||
|
||||
for propName, propValue := range properties {
|
||||
propMap, ok := propValue.(map[string]interface{})
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
|
||||
description := host.getPropertyDescription(propMap)
|
||||
schemaParams[propName] = description
|
||||
}
|
||||
|
||||
return host.marshalToJSON(schemaParams, "schema parameters")
|
||||
}
|
||||
|
||||
// getPropertyDescription extracts description from property map
|
||||
func (host *MCPHost) getPropertyDescription(propMap map[string]interface{}) string {
|
||||
if desc, exists := propMap["description"]; exists {
|
||||
if descStr, ok := desc.(string); ok {
|
||||
return descStr
|
||||
}
|
||||
}
|
||||
|
||||
// Fallback to type information
|
||||
if propType, exists := propMap["type"]; exists {
|
||||
if typeStr, ok := propType.(string); ok {
|
||||
return fmt.Sprintf("Parameter of type %s", typeStr)
|
||||
}
|
||||
}
|
||||
|
||||
return "Parameter"
|
||||
}
|
||||
|
||||
// extractReturns extracts return value information from ActionTool or docstring
|
||||
func (host *MCPHost) extractReturns(serverName, toolName string, info DocStringInfo) string {
|
||||
// Priority 1: Get from ActionTool interface if available
|
||||
if actionToolProvider := host.getActionToolProvider(serverName); actionToolProvider != nil {
|
||||
if actionTool := actionToolProvider.GetToolByAction(option.ActionName(toolName)); actionTool != nil {
|
||||
returnSchema := actionTool.ReturnSchema()
|
||||
if len(returnSchema) > 0 {
|
||||
return host.marshalToJSON(returnSchema, "return schema")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Priority 2: Use docstring returns as fallback
|
||||
if len(info.Returns) > 0 {
|
||||
return host.marshalToJSON(info.Returns, "docstring returns")
|
||||
}
|
||||
|
||||
return "{}"
|
||||
}
|
||||
|
||||
// marshalToJSON marshals data to JSON string with error handling
|
||||
func (host *MCPHost) marshalToJSON(data interface{}, dataType string) string {
|
||||
jsonBytes, err := sonic.MarshalString(data)
|
||||
if err != nil {
|
||||
log.Warn().Interface("data", data).Err(err).
|
||||
Msgf("failed to marshal %s to JSON", dataType)
|
||||
return "{}"
|
||||
}
|
||||
return jsonBytes
|
||||
}
|
||||
|
||||
// ExportToolsToJSON dumps MCP tools to JSON file
|
||||
func (h *MCPHost) ExportToolsToJSON(ctx context.Context, dumpPath string) error {
|
||||
// get all tools
|
||||
tools := h.GetTools(ctx)
|
||||
// convert to records
|
||||
records := ConvertToolsToRecords(tools)
|
||||
records := h.ConvertToolsToRecords(tools)
|
||||
// convert to JSON
|
||||
recordsJSON, err := sonic.MarshalIndent(records, "", " ")
|
||||
if err != nil {
|
||||
|
||||
Reference in New Issue
Block a user