mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-07-09 14:08:25 +02:00
Add files via upload
This commit is contained in:
+222
-110
@@ -193,12 +193,40 @@ type AgentLoopResult struct {
|
||||
MCPExecutionIDs []string
|
||||
}
|
||||
|
||||
// ProgressCallback 进度回调函数类型
|
||||
type ProgressCallback func(eventType, message string, data interface{})
|
||||
|
||||
// AgentLoop 执行Agent循环
|
||||
func (a *Agent) AgentLoop(ctx context.Context, userInput string, historyMessages []ChatMessage) (*AgentLoopResult, error) {
|
||||
return a.AgentLoopWithProgress(ctx, userInput, historyMessages, nil)
|
||||
}
|
||||
|
||||
// AgentLoopWithProgress 执行Agent循环(带进度回调)
|
||||
func (a *Agent) AgentLoopWithProgress(ctx context.Context, userInput string, historyMessages []ChatMessage, callback ProgressCallback) (*AgentLoopResult, error) {
|
||||
// 发送进度更新
|
||||
sendProgress := func(eventType, message string, data interface{}) {
|
||||
if callback != nil {
|
||||
callback(eventType, message, data)
|
||||
}
|
||||
}
|
||||
|
||||
// 系统提示词,指导AI如何处理工具错误
|
||||
systemPrompt := `你是一个专业的网络安全渗透测试专家。你可以使用各种安全工具进行自主渗透测试。分析目标并选择最佳测试策略。
|
||||
|
||||
重要:当工具调用失败时,请遵循以下原则:
|
||||
1. 仔细分析错误信息,理解失败的具体原因
|
||||
2. 如果工具不存在或未启用,尝试使用其他替代工具完成相同目标
|
||||
3. 如果参数错误,根据错误提示修正参数后重试
|
||||
4. 如果工具执行失败但输出了有用信息,可以基于这些信息继续分析
|
||||
5. 如果确实无法使用某个工具,向用户说明问题,并建议替代方案或手动操作
|
||||
6. 不要因为单个工具失败就停止整个测试流程,尝试其他方法继续完成任务
|
||||
|
||||
当工具返回错误时,错误信息会包含在工具响应中,请仔细阅读并做出合理的决策。`
|
||||
|
||||
messages := []ChatMessage{
|
||||
{
|
||||
Role: "system",
|
||||
Content: "你是一个专业的网络安全渗透测试专家。你可以使用各种安全工具进行自主渗透测试。分析目标并选择最佳测试策略。当需要执行工具时,使用提供的工具函数。",
|
||||
Content: systemPrompt,
|
||||
},
|
||||
}
|
||||
|
||||
@@ -248,6 +276,13 @@ func (a *Agent) AgentLoop(ctx context.Context, userInput string, historyMessages
|
||||
// 获取可用工具
|
||||
tools := a.getAvailableTools()
|
||||
|
||||
// 发送进度更新
|
||||
if i == 0 {
|
||||
sendProgress("progress", "正在分析请求并制定测试策略...", nil)
|
||||
} else {
|
||||
sendProgress("progress", fmt.Sprintf("正在继续分析(第 %d 轮迭代)...", i+1), nil)
|
||||
}
|
||||
|
||||
// 记录每次调用OpenAI
|
||||
if i == 0 {
|
||||
a.logger.Info("调用OpenAI",
|
||||
@@ -277,6 +312,7 @@ func (a *Agent) AgentLoop(ctx context.Context, userInput string, historyMessages
|
||||
}
|
||||
|
||||
// 调用OpenAI
|
||||
sendProgress("progress", "正在调用AI模型...", nil)
|
||||
response, err := a.callOpenAI(ctx, messages, tools)
|
||||
if err != nil {
|
||||
result.Response = ""
|
||||
@@ -304,17 +340,46 @@ func (a *Agent) AgentLoop(ctx context.Context, userInput string, historyMessages
|
||||
ToolCalls: choice.Message.ToolCalls,
|
||||
})
|
||||
|
||||
// 发送工具调用进度
|
||||
sendProgress("progress", fmt.Sprintf("检测到 %d 个工具调用,开始执行...", len(choice.Message.ToolCalls)), nil)
|
||||
|
||||
// 执行所有工具调用
|
||||
for _, toolCall := range choice.Message.ToolCalls {
|
||||
for idx, toolCall := range choice.Message.ToolCalls {
|
||||
// 发送工具调用开始事件
|
||||
toolArgsJSON, _ := json.Marshal(toolCall.Function.Arguments)
|
||||
sendProgress("tool_call", fmt.Sprintf("正在调用工具: %s", toolCall.Function.Name), map[string]interface{}{
|
||||
"toolName": toolCall.Function.Name,
|
||||
"arguments": string(toolArgsJSON),
|
||||
"index": idx + 1,
|
||||
"total": len(choice.Message.ToolCalls),
|
||||
})
|
||||
|
||||
// 执行工具
|
||||
execResult, err := a.executeToolViaMCP(ctx, toolCall.Function.Name, toolCall.Function.Arguments)
|
||||
if err != nil {
|
||||
// 构建详细的错误信息,帮助AI理解问题并做出决策
|
||||
errorMsg := a.formatToolError(toolCall.Function.Name, toolCall.Function.Arguments, err)
|
||||
messages = append(messages, ChatMessage{
|
||||
Role: "tool",
|
||||
ToolCallID: toolCall.ID,
|
||||
Content: fmt.Sprintf("工具执行失败: %v", err),
|
||||
Content: errorMsg,
|
||||
})
|
||||
|
||||
// 发送工具执行失败事件
|
||||
sendProgress("tool_result", fmt.Sprintf("工具 %s 执行失败", toolCall.Function.Name), map[string]interface{}{
|
||||
"toolName": toolCall.Function.Name,
|
||||
"success": false,
|
||||
"error": err.Error(),
|
||||
"index": idx + 1,
|
||||
"total": len(choice.Message.ToolCalls),
|
||||
})
|
||||
|
||||
a.logger.Warn("工具执行失败,已返回详细错误信息",
|
||||
zap.String("tool", toolCall.Function.Name),
|
||||
zap.Error(err),
|
||||
)
|
||||
} else {
|
||||
// 即使工具返回了错误结果(IsError=true),也继续处理,让AI决定下一步
|
||||
messages = append(messages, ChatMessage{
|
||||
Role: "tool",
|
||||
ToolCallID: toolCall.ID,
|
||||
@@ -324,6 +389,29 @@ func (a *Agent) AgentLoop(ctx context.Context, userInput string, historyMessages
|
||||
if execResult.ExecutionID != "" {
|
||||
result.MCPExecutionIDs = append(result.MCPExecutionIDs, execResult.ExecutionID)
|
||||
}
|
||||
|
||||
// 发送工具执行成功事件
|
||||
resultPreview := execResult.Result
|
||||
if len(resultPreview) > 200 {
|
||||
resultPreview = resultPreview[:200] + "..."
|
||||
}
|
||||
sendProgress("tool_result", fmt.Sprintf("工具 %s 执行完成", toolCall.Function.Name), map[string]interface{}{
|
||||
"toolName": toolCall.Function.Name,
|
||||
"success": !execResult.IsError,
|
||||
"isError": execResult.IsError,
|
||||
"result": resultPreview,
|
||||
"executionId": execResult.ExecutionID,
|
||||
"index": idx + 1,
|
||||
"total": len(choice.Message.ToolCalls),
|
||||
})
|
||||
|
||||
// 如果工具返回了错误,记录日志但不中断流程
|
||||
if execResult.IsError {
|
||||
a.logger.Warn("工具返回错误结果,但继续处理",
|
||||
zap.String("tool", toolCall.Function.Name),
|
||||
zap.String("result", execResult.Result),
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
continue
|
||||
@@ -337,6 +425,7 @@ func (a *Agent) AgentLoop(ctx context.Context, userInput string, historyMessages
|
||||
|
||||
// 如果完成,返回结果
|
||||
if choice.FinishReason == "stop" {
|
||||
sendProgress("progress", "正在生成最终回复...", nil)
|
||||
result.Response = choice.Message.Content
|
||||
return result, nil
|
||||
}
|
||||
@@ -347,117 +436,98 @@ func (a *Agent) AgentLoop(ctx context.Context, userInput string, historyMessages
|
||||
}
|
||||
|
||||
// getAvailableTools 获取可用工具
|
||||
// 从MCP服务器动态获取工具列表,使用简短描述以减少token消耗
|
||||
func (a *Agent) getAvailableTools() []Tool {
|
||||
// 从MCP服务器获取工具列表
|
||||
executions := a.mcpServer.GetAllExecutions()
|
||||
toolNames := make(map[string]bool)
|
||||
for _, exec := range executions {
|
||||
toolNames[exec.ToolName] = true
|
||||
// 从MCP服务器获取所有已注册的工具
|
||||
mcpTools := a.mcpServer.GetAllTools()
|
||||
|
||||
// 转换为OpenAI格式的工具定义
|
||||
tools := make([]Tool, 0, len(mcpTools))
|
||||
for _, mcpTool := range mcpTools {
|
||||
// 使用简短描述(如果存在),否则使用详细描述
|
||||
description := mcpTool.ShortDescription
|
||||
if description == "" {
|
||||
description = mcpTool.Description
|
||||
}
|
||||
|
||||
// 转换schema中的类型为OpenAI标准类型
|
||||
convertedSchema := a.convertSchemaTypes(mcpTool.InputSchema)
|
||||
|
||||
tools = append(tools, Tool{
|
||||
Type: "function",
|
||||
Function: FunctionDefinition{
|
||||
Name: mcpTool.Name,
|
||||
Description: description, // 使用简短描述减少token消耗
|
||||
Parameters: convertedSchema,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
tools := []Tool{
|
||||
{
|
||||
Type: "function",
|
||||
Function: FunctionDefinition{
|
||||
Name: "nmap",
|
||||
Description: "使用nmap进行网络扫描,发现开放端口和服务。支持IP地址、域名或URL(会自动提取域名)。使用TCP连接扫描,不需要root权限。",
|
||||
Parameters: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"target": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "目标IP地址、域名或URL(如 https://example.com)。如果是URL,会自动提取域名部分。",
|
||||
},
|
||||
"ports": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "要扫描的端口范围,例如: 1-1000 或 80,443,8080。如果不指定,将扫描常用端口。",
|
||||
},
|
||||
},
|
||||
"required": []string{"target"},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
Type: "function",
|
||||
Function: FunctionDefinition{
|
||||
Name: "sqlmap",
|
||||
Description: "使用sqlmap检测SQL注入漏洞",
|
||||
Parameters: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"url": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "目标URL",
|
||||
},
|
||||
},
|
||||
"required": []string{"url"},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
Type: "function",
|
||||
Function: FunctionDefinition{
|
||||
Name: "nikto",
|
||||
Description: "使用nikto扫描Web服务器漏洞",
|
||||
Parameters: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"target": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "目标URL",
|
||||
},
|
||||
},
|
||||
"required": []string{"target"},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
Type: "function",
|
||||
Function: FunctionDefinition{
|
||||
Name: "dirb",
|
||||
Description: "使用dirb进行目录扫描",
|
||||
Parameters: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"url": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "目标URL",
|
||||
},
|
||||
},
|
||||
"required": []string{"url"},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
Type: "function",
|
||||
Function: FunctionDefinition{
|
||||
Name: "exec",
|
||||
Description: "执行系统命令(谨慎使用,仅用于必要的系统操作)",
|
||||
Parameters: map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"command": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "要执行的系统命令",
|
||||
},
|
||||
"shell": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "使用的shell(可选,默认为sh)",
|
||||
},
|
||||
"workdir": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "工作目录(可选)",
|
||||
},
|
||||
},
|
||||
"required": []string{"command"},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
a.logger.Debug("获取可用工具列表",
|
||||
zap.Int("count", len(tools)),
|
||||
)
|
||||
|
||||
return tools
|
||||
}
|
||||
|
||||
// convertSchemaTypes 递归转换schema中的类型为OpenAI标准类型
|
||||
func (a *Agent) convertSchemaTypes(schema map[string]interface{}) map[string]interface{} {
|
||||
if schema == nil {
|
||||
return schema
|
||||
}
|
||||
|
||||
// 创建新的schema副本
|
||||
converted := make(map[string]interface{})
|
||||
for k, v := range schema {
|
||||
converted[k] = v
|
||||
}
|
||||
|
||||
// 转换properties中的类型
|
||||
if properties, ok := converted["properties"].(map[string]interface{}); ok {
|
||||
convertedProperties := make(map[string]interface{})
|
||||
for propName, propValue := range properties {
|
||||
if prop, ok := propValue.(map[string]interface{}); ok {
|
||||
convertedProp := make(map[string]interface{})
|
||||
for pk, pv := range prop {
|
||||
if pk == "type" {
|
||||
// 转换类型
|
||||
if typeStr, ok := pv.(string); ok {
|
||||
convertedProp[pk] = a.convertToOpenAIType(typeStr)
|
||||
} else {
|
||||
convertedProp[pk] = pv
|
||||
}
|
||||
} else {
|
||||
convertedProp[pk] = pv
|
||||
}
|
||||
}
|
||||
convertedProperties[propName] = convertedProp
|
||||
} else {
|
||||
convertedProperties[propName] = propValue
|
||||
}
|
||||
}
|
||||
converted["properties"] = convertedProperties
|
||||
}
|
||||
|
||||
return converted
|
||||
}
|
||||
|
||||
// convertToOpenAIType 将配置中的类型转换为OpenAI/JSON Schema标准类型
|
||||
func (a *Agent) convertToOpenAIType(configType string) string {
|
||||
switch configType {
|
||||
case "bool":
|
||||
return "boolean"
|
||||
case "int", "integer":
|
||||
return "number"
|
||||
case "float", "double":
|
||||
return "number"
|
||||
case "string", "array", "object":
|
||||
return configType
|
||||
default:
|
||||
// 默认返回原类型
|
||||
return configType
|
||||
}
|
||||
}
|
||||
|
||||
// callOpenAI 调用OpenAI API
|
||||
func (a *Agent) callOpenAI(ctx context.Context, messages []ChatMessage, tools []Tool) (*OpenAIResponse, error) {
|
||||
reqBody := OpenAIRequest{
|
||||
@@ -546,9 +616,11 @@ func (a *Agent) parseToolCall(content string) (map[string]interface{}, error) {
|
||||
type ToolExecutionResult struct {
|
||||
Result string
|
||||
ExecutionID string
|
||||
IsError bool // 标记是否为错误结果
|
||||
}
|
||||
|
||||
// executeToolViaMCP 通过MCP执行工具
|
||||
// 即使工具执行失败,也返回结果而不是错误,让AI能够处理错误情况
|
||||
func (a *Agent) executeToolViaMCP(ctx context.Context, toolName string, args map[string]interface{}) (*ToolExecutionResult, error) {
|
||||
a.logger.Info("通过MCP执行工具",
|
||||
zap.String("tool", toolName),
|
||||
@@ -557,8 +629,30 @@ func (a *Agent) executeToolViaMCP(ctx context.Context, toolName string, args map
|
||||
|
||||
// 通过MCP服务器调用工具
|
||||
result, executionID, err := a.mcpServer.CallTool(ctx, toolName, args)
|
||||
|
||||
// 如果调用失败(如工具不存在),返回友好的错误信息而不是抛出异常
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("工具执行失败: %w", err)
|
||||
errorMsg := fmt.Sprintf(`工具调用失败
|
||||
|
||||
工具名称: %s
|
||||
错误类型: 系统错误
|
||||
错误详情: %v
|
||||
|
||||
可能的原因:
|
||||
- 工具 "%s" 不存在或未启用
|
||||
- 系统配置问题
|
||||
- 网络或权限问题
|
||||
|
||||
建议:
|
||||
- 检查工具名称是否正确
|
||||
- 尝试使用其他替代工具
|
||||
- 如果这是必需的工具,请向用户说明情况`, toolName, err, toolName)
|
||||
|
||||
return &ToolExecutionResult{
|
||||
Result: errorMsg,
|
||||
ExecutionID: executionID,
|
||||
IsError: true,
|
||||
}, nil // 返回 nil 错误,让调用者处理结果
|
||||
}
|
||||
|
||||
// 格式化结果
|
||||
@@ -571,6 +665,24 @@ func (a *Agent) executeToolViaMCP(ctx context.Context, toolName string, args map
|
||||
return &ToolExecutionResult{
|
||||
Result: resultText.String(),
|
||||
ExecutionID: executionID,
|
||||
IsError: result != nil && result.IsError,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// formatToolError 格式化工具错误信息,提供更友好的错误描述
|
||||
func (a *Agent) formatToolError(toolName string, args map[string]interface{}, err error) string {
|
||||
errorMsg := fmt.Sprintf(`工具执行失败
|
||||
|
||||
工具名称: %s
|
||||
调用参数: %v
|
||||
错误信息: %v
|
||||
|
||||
请分析错误原因并采取以下行动之一:
|
||||
1. 如果参数错误,请修正参数后重试
|
||||
2. 如果工具不可用,请尝试使用替代工具
|
||||
3. 如果这是系统问题,请向用户说明情况并提供建议
|
||||
4. 如果错误信息中包含有用信息,可以基于这些信息继续分析`, toolName, args, err)
|
||||
|
||||
return errorMsg
|
||||
}
|
||||
|
||||
|
||||
@@ -115,6 +115,8 @@ func setupRoutes(router *gin.Engine, agentHandler *handler.AgentHandler, monitor
|
||||
{
|
||||
// Agent Loop
|
||||
api.POST("/agent-loop", agentHandler.AgentLoop)
|
||||
// Agent Loop 流式输出
|
||||
api.POST("/agent-loop/stream", agentHandler.AgentLoopStream)
|
||||
|
||||
// 对话历史
|
||||
api.POST("/conversations", conversationHandler.CreateConversation)
|
||||
|
||||
+107
-9
@@ -3,6 +3,8 @@ package config
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
|
||||
"gopkg.in/yaml.v3"
|
||||
)
|
||||
@@ -39,7 +41,8 @@ type OpenAIConfig struct {
|
||||
}
|
||||
|
||||
type SecurityConfig struct {
|
||||
Tools []ToolConfig `yaml:"tools"`
|
||||
Tools []ToolConfig `yaml:"tools,omitempty"` // 向后兼容:支持在主配置文件中定义工具
|
||||
ToolsDir string `yaml:"tools_dir,omitempty"` // 工具配置文件目录(新方式)
|
||||
}
|
||||
|
||||
type DatabaseConfig struct {
|
||||
@@ -47,13 +50,14 @@ type DatabaseConfig struct {
|
||||
}
|
||||
|
||||
type ToolConfig struct {
|
||||
Name string `yaml:"name"`
|
||||
Command string `yaml:"command"`
|
||||
Args []string `yaml:"args,omitempty"` // 固定参数(可选)
|
||||
Description string `yaml:"description"`
|
||||
Enabled bool `yaml:"enabled"`
|
||||
Parameters []ParameterConfig `yaml:"parameters,omitempty"` // 参数定义(可选)
|
||||
ArgMapping string `yaml:"arg_mapping,omitempty"` // 参数映射方式: "auto", "manual", "template"(可选)
|
||||
Name string `yaml:"name"`
|
||||
Command string `yaml:"command"`
|
||||
Args []string `yaml:"args,omitempty"` // 固定参数(可选)
|
||||
ShortDescription string `yaml:"short_description,omitempty"` // 简短描述(用于工具列表,减少token消耗)
|
||||
Description string `yaml:"description"` // 详细描述(用于工具文档)
|
||||
Enabled bool `yaml:"enabled"`
|
||||
Parameters []ParameterConfig `yaml:"parameters,omitempty"` // 参数定义(可选)
|
||||
ArgMapping string `yaml:"arg_mapping,omitempty"` // 参数映射方式: "auto", "manual", "template"(可选)
|
||||
}
|
||||
|
||||
// ParameterConfig 参数配置
|
||||
@@ -81,9 +85,102 @@ func Load(path string) (*Config, error) {
|
||||
return nil, fmt.Errorf("解析配置文件失败: %w", err)
|
||||
}
|
||||
|
||||
// 如果配置了工具目录,从目录加载工具配置
|
||||
if cfg.Security.ToolsDir != "" {
|
||||
configDir := filepath.Dir(path)
|
||||
toolsDir := cfg.Security.ToolsDir
|
||||
|
||||
// 如果是相对路径,相对于配置文件所在目录
|
||||
if !filepath.IsAbs(toolsDir) {
|
||||
toolsDir = filepath.Join(configDir, toolsDir)
|
||||
}
|
||||
|
||||
tools, err := LoadToolsFromDir(toolsDir)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("从工具目录加载工具配置失败: %w", err)
|
||||
}
|
||||
|
||||
// 合并工具配置:目录中的工具优先,主配置中的工具作为补充
|
||||
existingTools := make(map[string]bool)
|
||||
for _, tool := range tools {
|
||||
existingTools[tool.Name] = true
|
||||
}
|
||||
|
||||
// 添加主配置中不存在于目录中的工具(向后兼容)
|
||||
for _, tool := range cfg.Security.Tools {
|
||||
if !existingTools[tool.Name] {
|
||||
tools = append(tools, tool)
|
||||
}
|
||||
}
|
||||
|
||||
cfg.Security.Tools = tools
|
||||
}
|
||||
|
||||
return &cfg, nil
|
||||
}
|
||||
|
||||
// LoadToolsFromDir 从目录加载所有工具配置文件
|
||||
func LoadToolsFromDir(dir string) ([]ToolConfig, error) {
|
||||
var tools []ToolConfig
|
||||
|
||||
// 检查目录是否存在
|
||||
if _, err := os.Stat(dir); os.IsNotExist(err) {
|
||||
return tools, nil // 目录不存在时返回空列表,不报错
|
||||
}
|
||||
|
||||
// 读取目录中的所有 .yaml 和 .yml 文件
|
||||
entries, err := os.ReadDir(dir)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("读取工具目录失败: %w", err)
|
||||
}
|
||||
|
||||
for _, entry := range entries {
|
||||
if entry.IsDir() {
|
||||
continue
|
||||
}
|
||||
|
||||
name := entry.Name()
|
||||
if !strings.HasSuffix(name, ".yaml") && !strings.HasSuffix(name, ".yml") {
|
||||
continue
|
||||
}
|
||||
|
||||
filePath := filepath.Join(dir, name)
|
||||
tool, err := LoadToolFromFile(filePath)
|
||||
if err != nil {
|
||||
// 记录错误但继续加载其他文件
|
||||
fmt.Printf("警告: 加载工具配置文件 %s 失败: %v\n", filePath, err)
|
||||
continue
|
||||
}
|
||||
|
||||
tools = append(tools, *tool)
|
||||
}
|
||||
|
||||
return tools, nil
|
||||
}
|
||||
|
||||
// LoadToolFromFile 从单个文件加载工具配置
|
||||
func LoadToolFromFile(path string) (*ToolConfig, error) {
|
||||
data, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("读取文件失败: %w", err)
|
||||
}
|
||||
|
||||
var tool ToolConfig
|
||||
if err := yaml.Unmarshal(data, &tool); err != nil {
|
||||
return nil, fmt.Errorf("解析工具配置失败: %w", err)
|
||||
}
|
||||
|
||||
// 验证必需字段
|
||||
if tool.Name == "" {
|
||||
return nil, fmt.Errorf("工具名称不能为空")
|
||||
}
|
||||
if tool.Command == "" {
|
||||
return nil, fmt.Errorf("工具命令不能为空")
|
||||
}
|
||||
|
||||
return &tool, nil
|
||||
}
|
||||
|
||||
func Default() *Config {
|
||||
return &Config{
|
||||
Server: ServerConfig{
|
||||
@@ -104,7 +201,8 @@ func Default() *Config {
|
||||
Model: "gpt-4",
|
||||
},
|
||||
Security: SecurityConfig{
|
||||
Tools: []ToolConfig{}, // 工具配置应该从 config.yaml 加载,不在此硬编码
|
||||
Tools: []ToolConfig{}, // 工具配置应该从 config.yaml 或 tools/ 目录加载
|
||||
ToolsDir: "tools", // 默认工具目录
|
||||
},
|
||||
Database: DatabaseConfig{
|
||||
Path: "data/conversations.db",
|
||||
|
||||
@@ -1,6 +1,8 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
@@ -132,3 +134,120 @@ func (h *AgentHandler) AgentLoop(c *gin.Context) {
|
||||
})
|
||||
}
|
||||
|
||||
// StreamEvent 流式事件
|
||||
type StreamEvent struct {
|
||||
Type string `json:"type"` // progress, tool_call, tool_result, response, error, done
|
||||
Message string `json:"message"` // 显示消息
|
||||
Data interface{} `json:"data,omitempty"`
|
||||
}
|
||||
|
||||
// AgentLoopStream 处理Agent Loop流式请求
|
||||
func (h *AgentHandler) AgentLoopStream(c *gin.Context) {
|
||||
var req ChatRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
// 对于流式请求,也发送SSE格式的错误
|
||||
c.Header("Content-Type", "text/event-stream")
|
||||
c.Header("Cache-Control", "no-cache")
|
||||
c.Header("Connection", "keep-alive")
|
||||
event := StreamEvent{
|
||||
Type: "error",
|
||||
Message: "请求参数错误: " + err.Error(),
|
||||
}
|
||||
eventJSON, _ := json.Marshal(event)
|
||||
fmt.Fprintf(c.Writer, "data: %s\n\n", eventJSON)
|
||||
c.Writer.Flush()
|
||||
return
|
||||
}
|
||||
|
||||
h.logger.Info("收到Agent Loop流式请求",
|
||||
zap.String("message", req.Message),
|
||||
zap.String("conversationId", req.ConversationID),
|
||||
)
|
||||
|
||||
// 设置SSE响应头
|
||||
c.Header("Content-Type", "text/event-stream")
|
||||
c.Header("Cache-Control", "no-cache")
|
||||
c.Header("Connection", "keep-alive")
|
||||
c.Header("X-Accel-Buffering", "no") // 禁用nginx缓冲
|
||||
|
||||
// 发送初始事件
|
||||
sendEvent := func(eventType, message string, data interface{}) {
|
||||
event := StreamEvent{
|
||||
Type: eventType,
|
||||
Message: message,
|
||||
Data: data,
|
||||
}
|
||||
eventJSON, _ := json.Marshal(event)
|
||||
fmt.Fprintf(c.Writer, "data: %s\n\n", eventJSON)
|
||||
c.Writer.Flush()
|
||||
}
|
||||
|
||||
// 如果没有对话ID,创建新对话
|
||||
conversationID := req.ConversationID
|
||||
if conversationID == "" {
|
||||
title := req.Message
|
||||
if len(title) > 50 {
|
||||
title = title[:50] + "..."
|
||||
}
|
||||
conv, err := h.db.CreateConversation(title)
|
||||
if err != nil {
|
||||
h.logger.Error("创建对话失败", zap.Error(err))
|
||||
sendEvent("error", "创建对话失败: "+err.Error(), nil)
|
||||
return
|
||||
}
|
||||
conversationID = conv.ID
|
||||
}
|
||||
|
||||
// 获取历史消息
|
||||
historyMessages, err := h.db.GetMessages(conversationID)
|
||||
if err != nil {
|
||||
h.logger.Warn("获取历史消息失败", zap.Error(err))
|
||||
historyMessages = []database.Message{}
|
||||
}
|
||||
|
||||
// 将数据库消息转换为Agent消息格式
|
||||
agentHistoryMessages := make([]agent.ChatMessage, 0, len(historyMessages))
|
||||
for _, msg := range historyMessages {
|
||||
agentHistoryMessages = append(agentHistoryMessages, agent.ChatMessage{
|
||||
Role: msg.Role,
|
||||
Content: msg.Content,
|
||||
})
|
||||
}
|
||||
|
||||
// 保存用户消息
|
||||
_, err = h.db.AddMessage(conversationID, "user", req.Message, nil)
|
||||
if err != nil {
|
||||
h.logger.Error("保存用户消息失败", zap.Error(err))
|
||||
}
|
||||
|
||||
// 创建进度回调函数
|
||||
progressCallback := func(eventType, message string, data interface{}) {
|
||||
sendEvent(eventType, message, data)
|
||||
}
|
||||
|
||||
// 执行Agent Loop,传入进度回调
|
||||
sendEvent("progress", "正在分析您的请求...", nil)
|
||||
result, err := h.agent.AgentLoopWithProgress(c.Request.Context(), req.Message, agentHistoryMessages, progressCallback)
|
||||
if err != nil {
|
||||
h.logger.Error("Agent Loop执行失败", zap.Error(err))
|
||||
sendEvent("error", "执行失败: "+err.Error(), nil)
|
||||
sendEvent("done", "", nil)
|
||||
return
|
||||
}
|
||||
|
||||
// 保存助手回复
|
||||
_, err = h.db.AddMessage(conversationID, "assistant", result.Response, result.MCPExecutionIDs)
|
||||
if err != nil {
|
||||
h.logger.Error("保存助手消息失败", zap.Error(err))
|
||||
}
|
||||
|
||||
// 发送最终响应
|
||||
sendEvent("response", result.Response, map[string]interface{}{
|
||||
"mcpExecutionIds": result.MCPExecutionIDs,
|
||||
"conversationId": conversationID,
|
||||
})
|
||||
sendEvent("done", "", map[string]interface{}{
|
||||
"conversationId": conversationID,
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
@@ -361,6 +361,18 @@ func (s *Server) GetStats() map[string]*ToolStats {
|
||||
return stats
|
||||
}
|
||||
|
||||
// GetAllTools 获取所有已注册的工具(用于Agent动态获取工具列表)
|
||||
func (s *Server) GetAllTools() []Tool {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
|
||||
tools := make([]Tool, 0, len(s.toolDefs))
|
||||
for _, tool := range s.toolDefs {
|
||||
tools = append(tools, tool)
|
||||
}
|
||||
return tools
|
||||
}
|
||||
|
||||
// CallTool 直接调用工具(用于内部调用)
|
||||
func (s *Server) CallTool(ctx context.Context, toolName string, args map[string]interface{}) (*ToolResult, string, error) {
|
||||
s.mu.RLock()
|
||||
|
||||
@@ -36,9 +36,10 @@ type Error struct {
|
||||
|
||||
// Tool 表示MCP工具定义
|
||||
type Tool struct {
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description"`
|
||||
InputSchema map[string]interface{} `json:"inputSchema"`
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description"` // 详细描述
|
||||
ShortDescription string `json:"shortDescription,omitempty"` // 简短描述(用于工具列表,减少token消耗)
|
||||
InputSchema map[string]interface{} `json:"inputSchema"`
|
||||
}
|
||||
|
||||
// ToolCall 表示工具调用
|
||||
|
||||
@@ -150,10 +150,28 @@ func (e *Executor) RegisterTools(mcpServer *mcp.Server) {
|
||||
toolName := toolConfig.Name
|
||||
toolConfigCopy := toolConfig
|
||||
|
||||
// 使用简短描述(如果存在),否则使用详细描述的前100个字符
|
||||
shortDesc := toolConfigCopy.ShortDescription
|
||||
if shortDesc == "" {
|
||||
// 如果没有简短描述,从详细描述中提取第一行或前100个字符
|
||||
desc := toolConfigCopy.Description
|
||||
if len(desc) > 100 {
|
||||
// 尝试找到第一个换行符
|
||||
if idx := strings.Index(desc, "\n"); idx > 0 && idx < 100 {
|
||||
shortDesc = strings.TrimSpace(desc[:idx])
|
||||
} else {
|
||||
shortDesc = desc[:100] + "..."
|
||||
}
|
||||
} else {
|
||||
shortDesc = desc
|
||||
}
|
||||
}
|
||||
|
||||
tool := mcp.Tool{
|
||||
Name: toolConfigCopy.Name,
|
||||
Description: toolConfigCopy.Description,
|
||||
InputSchema: e.buildInputSchema(&toolConfigCopy),
|
||||
Name: toolConfigCopy.Name,
|
||||
Description: toolConfigCopy.Description,
|
||||
ShortDescription: shortDesc,
|
||||
InputSchema: e.buildInputSchema(&toolConfigCopy),
|
||||
}
|
||||
|
||||
handler := func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||
@@ -543,8 +561,11 @@ func (e *Executor) buildInputSchema(toolConfig *config.ToolConfig) map[string]in
|
||||
required := []string{}
|
||||
|
||||
for _, param := range toolConfig.Parameters {
|
||||
// 转换类型为OpenAI/JSON Schema标准类型
|
||||
openAIType := e.convertToOpenAIType(param.Type)
|
||||
|
||||
prop := map[string]interface{}{
|
||||
"type": param.Type,
|
||||
"type": openAIType,
|
||||
"description": param.Description,
|
||||
}
|
||||
|
||||
@@ -622,6 +643,26 @@ func (e *Executor) buildInputSchema(toolConfig *config.ToolConfig) map[string]in
|
||||
return schema
|
||||
}
|
||||
|
||||
// convertToOpenAIType 将配置中的类型转换为OpenAI/JSON Schema标准类型
|
||||
func (e *Executor) convertToOpenAIType(configType string) string {
|
||||
switch configType {
|
||||
case "bool":
|
||||
return "boolean"
|
||||
case "int", "integer":
|
||||
return "number"
|
||||
case "float", "double":
|
||||
return "number"
|
||||
case "string", "array", "object":
|
||||
return configType
|
||||
default:
|
||||
// 默认返回原类型,但记录警告
|
||||
e.logger.Warn("未知的参数类型,使用原类型",
|
||||
zap.String("type", configType),
|
||||
)
|
||||
return configType
|
||||
}
|
||||
}
|
||||
|
||||
// Vulnerability 漏洞信息
|
||||
type Vulnerability struct {
|
||||
ID string `json:"id"`
|
||||
|
||||
Reference in New Issue
Block a user