190 lines
4.8 KiB
Go
190 lines
4.8 KiB
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
|
|
}
|