Golang开发LLM工具调用验证SDK的经验分享

在Go语言中开发LLM工具调用验证SDK确实是个重要课题。hguard-go项目提供了一个很好的起点，但实际应用中需要更细粒度的控制。下面是一个基于策略验证的示例实现：

package main

import (
	"context"
	"fmt"
	"regexp"
	"strings"
)

// 工具调用验证器接口
type ToolCallValidator interface {
	Validate(ctx context.Context, call ToolCall) (bool, error)
}

// 工具调用结构
type ToolCall struct {
	ToolName    string
	Parameters  map[string]interface{}
	Source      string // 来源标识，如"claude", "gpt-4"
}

// 基于正则表达式的路径验证器
type PathValidator struct {
	allowedPatterns []*regexp.Regexp
	deniedPatterns  []*regexp.Regexp
}

func NewPathValidator(allowed, denied []string) (*PathValidator, error) {
	pv := &PathValidator{}
	
	for _, pattern := range allowed {
		re, err := regexp.Compile(pattern)
		if err != nil {
			return nil, err
		}
		pv.allowedPatterns = append(pv.allowedPatterns, re)
	}
	
	for _, pattern := range denied {
		re, err := regexp.Compile(pattern)
		if err != nil {
			return nil, err
		}
		pv.deniedPatterns = append(pv.deniedPatterns, re)
	}
	
	return pv, nil
}

func (pv *PathValidator) ValidatePath(path string) bool {
	// 先检查拒绝列表
	for _, pattern := range pv.deniedPatterns {
		if pattern.MatchString(path) {
			return false
		}
	}
	
	// 再检查允许列表
	for _, pattern := range pv.allowedPatterns {
		if pattern.MatchString(path) {
			return true
		}
	}
	
	return len(pv.allowedPatterns) == 0 // 如果没设置允许列表，默认拒绝
}

// 文件系统工具验证器
type FileSystemValidator struct {
	pathValidator *PathValidator
	maxFileSize   int64
	allowedOps    map[string]bool
}

func NewFileSystemValidator() *FileSystemValidator {
	return &FileSystemValidator{
		allowedOps: map[string]bool{
			"read":   true,
			"list":   true,
			"stat":   true,
			"exists": true,
		},
		maxFileSize: 10 * 1024 * 1024, // 10MB限制
	}
}

func (fv *FileSystemValidator) Validate(ctx context.Context, call ToolCall) (bool, error) {
	if call.ToolName != "filesystem" {
		return false, fmt.Errorf("invalid tool name")
	}
	
	// 验证操作类型
	op, ok := call.Parameters["operation"].(string)
	if !ok {
		return false, fmt.Errorf("missing operation parameter")
	}
	
	if !fv.allowedOps[op] {
		return false, fmt.Errorf("operation %s not allowed", op)
	}
	
	// 验证路径
	if path, ok := call.Parameters["path"].(string); ok {
		if !fv.pathValidator.ValidatePath(path) {
			return false, fmt.Errorf("path %s not allowed", path)
		}
	}
	
	// 验证文件大小限制
	if size, ok := call.Parameters["size"].(int64); ok && size > fv.maxFileSize {
		return false, fmt.Errorf("file size %d exceeds limit %d", size, fv.maxFileSize)
	}
	
	return true, nil
}

// 数据库查询验证器
type DatabaseValidator struct {
	allowedTables    map[string]bool
	maxRowsReturned  int
	readOnly         bool
}

func NewDatabaseValidator(allowedTables []string, readOnly bool) *DatabaseValidator {
	dv := &DatabaseValidator{
		allowedTables:   make(map[string]bool),
		maxRowsReturned: 1000,
		readOnly:        readOnly,
	}
	
	for _, table := range allowedTables {
		dv.allowedTables[table] = true
	}
	
	return dv
}

func (dv *DatabaseValidator) Validate(ctx context.Context, call ToolCall) (bool, error) {
	if call.ToolName != "database" {
		return false, fmt.Errorf("invalid tool name")
	}
	
	query, ok := call.Parameters["query"].(string)
	if !ok {
		return false, fmt.Errorf("missing query parameter")
	}
	
	// 检查是否为只读操作
	if dv.readOnly && dv.isWriteQuery(query) {
		return false, fmt.Errorf("write operations not allowed")
	}
	
	// 验证表访问权限
	tables := dv.extractTables(query)
	for _, table := range tables {
		if !dv.allowedTables[table] {
			return false, fmt.Errorf("access to table %s not allowed", table)
		}
	}
	
	// 验证返回行数限制
	if limit, ok := call.Parameters["limit"].(int); ok && limit > dv.maxRowsReturned {
		return false, fmt.Errorf("row limit %d exceeds maximum %d", limit, dv.maxRowsReturned)
	}
	
	return true, nil
}

func (dv *DatabaseValidator) isWriteQuery(query string) bool {
	query = strings.ToUpper(strings.TrimSpace(query))
	return strings.HasPrefix(query, "INSERT") ||
		strings.HasPrefix(query, "UPDATE") ||
		strings.HasPrefix(query, "DELETE") ||
		strings.HasPrefix(query, "DROP") ||
		strings.HasPrefix(query, "CREATE") ||
		strings.HasPrefix(query, "ALTER")
}

func (dv *DatabaseValidator) extractTables(query string) []string {
	// 简化的表名提取逻辑
	// 实际应用中应该使用SQL解析器
	re := regexp.MustCompile(`(?i)(?:FROM|JOIN|INTO|UPDATE)\s+(\w+)`)
	matches := re.FindAllStringSubmatch(query, -1)
	
	tables := make([]string, 0, len(matches))
	for _, match := range matches {
		if len(match) > 1 {
			tables = append(tables, strings.ToLower(match[1]))
		}
	}
	
	return tables
}

// 组合验证器
type CompositeValidator struct {
	validators map[string]ToolCallValidator
}

func NewCompositeValidator() *CompositeValidator {
	return &CompositeValidator{
		validators: make(map[string]ToolCallValidator),
	}
}

func (cv *CompositeValidator) RegisterValidator(toolName string, validator ToolCallValidator) {
	cv.validators[toolName] = validator
}

func (cv *CompositeValidator) Validate(ctx context.Context, call ToolCall) (bool, error) {
	validator, exists := cv.validators[call.ToolName]
	if !exists {
		return false, fmt.Errorf("no validator registered for tool %s", call.ToolName)
	}
	
	return validator.Validate(ctx, call)
}

// 使用示例
func main() {
	// 创建路径验证器
	pathValidator, _ := NewPathValidator(
		[]string{`^/tmp/.*$`, `^/home/user/documents/.*\.(txt|md|pdf)$`},
		[]string{`.*/\.\..*$`, `.*/etc/.*$`, `.*/root/.*$`},
	)
	
	// 创建文件系统验证器
	fsValidator := NewFileSystemValidator()
	fsValidator.pathValidator = pathValidator
	
	// 创建数据库验证器
	dbValidator := NewDatabaseValidator([]string{"users", "products", "orders"}, true)
	
	// 创建组合验证器
	compositeValidator := NewCompositeValidator()
	compositeValidator.RegisterValidator("filesystem", fsValidator)
	compositeValidator.RegisterValidator("database", dbValidator)
	
	// 测试验证
	ctx := context.Background()
	
	// 有效的文件读取调用
	validFileCall := ToolCall{
		ToolName: "filesystem",
		Parameters: map[string]interface{}{
			"operation": "read",
			"path":      "/tmp/test.txt",
		},
		Source: "claude",
	}
	
	valid, err := compositeValidator.Validate(ctx, validFileCall)
	fmt.Printf("Valid file call: %v, error: %v\n", valid, err)
	
	// 无效的文件写入调用
	invalidFileCall := ToolCall{
		ToolName: "filesystem",
		Parameters: map[string]interface{}{
			"operation": "write",
			"path":      "/etc/passwd",
		},
		Source: "claude",
	}
	
	valid, err = compositeValidator.Validate(ctx, invalidFileCall)
	fmt.Printf("Invalid file call: %v, error: %v\n", valid, err)
	
	// 有效的数据库查询
	validDBCall := ToolCall{
		ToolName: "database",
		Parameters: map[string]interface{}{
			"query": "SELECT * FROM users WHERE id = 1",
			"limit": 10,
		},
		Source: "claude",
	}
	
	valid, err = compositeValidator.Validate(ctx, validDBCall)
	fmt.Printf("Valid DB call: %v, error: %v\n", valid, err)
}

这个实现提供了几个关键特性：

1. 路径验证：使用正则表达式模式匹配来控制文件系统访问
2. 操作白名单：只允许特定的文件系统操作
3. 大小限制：防止读取或写入过大的文件
4. SQL查询分析：检测写操作和验证表访问权限
5. 组合验证：支持多种工具的集中验证

实际部署时，还需要考虑：

• 添加请求频率限制
• 实现审计日志记录所有工具调用
• 集成JWT或API密钥认证
• 添加请求超时控制
• 实现基于角色的访问控制(RBAC)

验证逻辑应该作为LLM工具调用的前置中间件，在工具执行前进行安全检查。这样可以确保即使LLM产生有害的工具调用，也能在到达实际系统前被拦截。