sdk_generate/internal/prompts/tools/crypt/common.go

// common.go
package crypt

import (
	"context"
	"encoding/json"
	"fmt"
	"strings"
	"sync"

	"github.com/sashabaranov/go-openai"
	"github.com/sashabaranov/go-openai/jsonschema"
)

// CryptoTool 接口 - 支持渐进式披露
type CryptoTool interface {
	Name() string
	Description() string
	GetFunctionDefinition() openai.FunctionDefinition
	GetDetail(ctx context.Context, params map[string]string) (string, error)
	Execute(ctx context.Context, params map[string]string) (string, error)
}

// ToolRegistry 工具注册中心
type ToolRegistry struct {
	tools map[string]CryptoTool
	mu    sync.RWMutex
}

func NewToolRegistry() *ToolRegistry {
	return &ToolRegistry{
		tools: map[string]CryptoTool{
			// 非对称加密
			"rsa_sign": &RSASignTool{},
			"sm2_sign": &SM2SignTool{},
			// 对称加密
			"aes_cbc_encrypt": &AESCBCEncryptTool{},
			"aes_ecb_encrypt": &AESECBEncryptTool{},
			"sm4_cbc_encrypt": &SM4CBCEncryptTool{},
			"sm4_ecb_encrypt": &SM4ECBEncryptTool{},
			// 哈希
			"sm3_hash": &SM3HashTool{},
			// 辅助工具
			"param_concat":    &ParamConcatTool{},
			"nonce_timestamp": &NonceTimestampTool{},
		},
	}
}

// GetToolManifest 获取所有工具的清单
func (r *ToolRegistry) GetToolManifest() []openai.Tool {
	r.mu.RLock()
	defer r.mu.RUnlock()

	tools := make([]openai.Tool, 0, len(r.tools))
	for _, tool := range r.tools {
		def := tool.GetFunctionDefinition()
		tools = append(tools, openai.Tool{
			Type: openai.ToolTypeFunction,
			Function: &openai.FunctionDefinition{
				Name:        def.Name,
				Description: def.Description,
				Parameters:  def.Parameters,
			},
		})
	}
	return tools
}

// GetToolDescriptions 获取所有工具的描述文本
func (r *ToolRegistry) GetToolDescriptions() string {
	r.mu.RLock()
	defer r.mu.RUnlock()

	var sb strings.Builder
	sb.WriteString("## 可用加密工具清单\n\n")
	sb.WriteString("当需要生成特定加密算法的代码时，请调用对应的工具获取完整实现。\n\n")

	// 按类别分组
	categories := map[string][]CryptoTool{
		"非对称加密": {},
		"对称加密":  {},
		"哈希算法":  {},
		"辅助工具":  {},
	}

	for _, tool := range r.tools {
		name := tool.Name()
		switch {
		case strings.HasPrefix(name, "rsa_") || strings.HasPrefix(name, "sm2_"):
			categories["非对称加密"] = append(categories["非对称加密"], tool)
		case strings.HasPrefix(name, "aes_") || strings.HasPrefix(name, "sm4_"):
			categories["对称加密"] = append(categories["对称加密"], tool)
		case strings.HasPrefix(name, "sm3_"):
			categories["哈希算法"] = append(categories["哈希算法"], tool)
		default:
			categories["辅助工具"] = append(categories["辅助工具"], tool)
		}
	}

	for category, tools := range categories {
		if len(tools) == 0 {
			continue
		}
		sb.WriteString(fmt.Sprintf("### %s\n\n", category))
		for _, tool := range tools {
			sb.WriteString(fmt.Sprintf("#### %s\n", tool.Name()))
			sb.WriteString(fmt.Sprintf("**描述**: %s\n", tool.Description()))

			def := tool.GetFunctionDefinition()
			if def.Parameters != nil {
				if schema, ok := def.Parameters.(jsonschema.Definition); ok && schema.Properties != nil {
					sb.WriteString("**参数**:\n")
					for name, prop := range schema.Properties {
						desc := prop.Description
						if desc == "" {
							desc = name
						}
						sb.WriteString(fmt.Sprintf("  - `%s`: %s\n", name, desc))
					}
				}
			}
			sb.WriteString("\n")
		}
	}

	return sb.String()
}

// GetToolDetail 获取工具详细实现
func (r *ToolRegistry) GetToolDetail(ctx context.Context, name string, params map[string]string) (string, error) {
	r.mu.RLock()
	tool, ok := r.tools[name]
	r.mu.RUnlock()

	if !ok {
		return "", fmt.Errorf("未知工具: %s", name)
	}
	return tool.GetDetail(ctx, params)
}

// ExecuteTool 执行工具
func (r *ToolRegistry) ExecuteTool(ctx context.Context, name string, params map[string]string) (string, error) {
	r.mu.RLock()
	tool, ok := r.tools[name]
	r.mu.RUnlock()

	if !ok {
		return "", fmt.Errorf("未知工具: %s", name)
	}
	return tool.Execute(ctx, params)
}

// ToolCallResult 工具调用结果结构
type ToolCallResult struct {
	ToolName    string `json:"tool_name"`
	Description string `json:"description"`
	Detail      string `json:"detail,omitempty"`
}

// BuildToolCallResult 构建工具调用结果
func (r *ToolRegistry) BuildToolCallResult(ctx context.Context, toolName string, params map[string]string, includeDetail bool) (string, error) {
	r.mu.RLock()
	tool, ok := r.tools[toolName]
	r.mu.RUnlock()

	if !ok {
		return "", fmt.Errorf("未知工具: %s", toolName)
	}

	resp := ToolCallResult{
		ToolName:    tool.Name(),
		Description: tool.Description(),
	}

	if includeDetail {
		detail, err := tool.GetDetail(ctx, params)
		if err != nil {
			return "", err
		}
		resp.Detail = detail
	}

	data, err := json.MarshalIndent(resp, "", "  ")
	if err != nil {
		return "", err
	}
	return string(data), nil
}