mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-08-08 20:06:07 +02:00
Add files via upload
This commit is contained in:
+30
-11
@@ -39,6 +39,7 @@ type Agent struct {
|
|||||||
toolNameMapping map[string]string // 工具名称映射:OpenAI格式 -> 原始格式(用于外部MCP工具)
|
toolNameMapping map[string]string // 工具名称映射:OpenAI格式 -> 原始格式(用于外部MCP工具)
|
||||||
currentConversationID string // 当前对话ID(用于自动传递给工具)
|
currentConversationID string // 当前对话ID(用于自动传递给工具)
|
||||||
promptBaseDir string // 解析 system_prompt_path 时相对路径的基准目录(通常为 config.yaml 所在目录)
|
promptBaseDir string // 解析 system_prompt_path 时相对路径的基准目录(通常为 config.yaml 所在目录)
|
||||||
|
toolDescriptionMode string // 工具描述模式: "short" | "full",默认 short
|
||||||
}
|
}
|
||||||
|
|
||||||
// ResultStorage 结果存储接口(直接使用 storage 包的类型)
|
// ResultStorage 结果存储接口(直接使用 storage 包的类型)
|
||||||
@@ -162,6 +163,7 @@ func NewAgent(cfg *config.OpenAIConfig, agentCfg *config.AgentConfig, mcpServer
|
|||||||
resultStorage: resultStorage,
|
resultStorage: resultStorage,
|
||||||
largeResultThreshold: largeResultThreshold,
|
largeResultThreshold: largeResultThreshold,
|
||||||
toolNameMapping: make(map[string]string), // 初始化工具名称映射
|
toolNameMapping: make(map[string]string), // 初始化工具名称映射
|
||||||
|
toolDescriptionMode: "short",
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -922,7 +924,7 @@ func (a *Agent) AgentLoopWithProgress(ctx context.Context, userInput string, his
|
|||||||
}
|
}
|
||||||
|
|
||||||
// getAvailableTools 获取可用工具
|
// getAvailableTools 获取可用工具
|
||||||
// 从MCP服务器动态获取工具列表,使用简短描述以减少token消耗
|
// 从MCP服务器动态获取工具列表,描述模式由 tool_description_mode 控制
|
||||||
// roleTools: 角色配置的工具列表(toolKey格式),如果为空或nil,则使用所有工具(默认角色)
|
// roleTools: 角色配置的工具列表(toolKey格式),如果为空或nil,则使用所有工具(默认角色)
|
||||||
func (a *Agent) getAvailableTools(roleTools []string) []Tool {
|
func (a *Agent) getAvailableTools(roleTools []string) []Tool {
|
||||||
// 构建角色工具集合(用于快速查找)
|
// 构建角色工具集合(用于快速查找)
|
||||||
@@ -946,11 +948,7 @@ func (a *Agent) getAvailableTools(roleTools []string) []Tool {
|
|||||||
continue // 不在角色工具列表中,跳过
|
continue // 不在角色工具列表中,跳过
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
// 使用简短描述(如果存在),否则使用详细描述
|
description := a.pickToolDescription(mcpTool.ShortDescription, mcpTool.Description)
|
||||||
description := mcpTool.ShortDescription
|
|
||||||
if description == "" {
|
|
||||||
description = mcpTool.Description
|
|
||||||
}
|
|
||||||
|
|
||||||
// 转换schema中的类型为OpenAI标准类型
|
// 转换schema中的类型为OpenAI标准类型
|
||||||
convertedSchema := a.convertSchemaTypes(mcpTool.InputSchema)
|
convertedSchema := a.convertSchemaTypes(mcpTool.InputSchema)
|
||||||
@@ -1024,11 +1022,7 @@ func (a *Agent) getAvailableTools(roleTools []string) []Tool {
|
|||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
|
|
||||||
// 使用简短描述(如果存在),否则使用详细描述
|
description := a.pickToolDescription(externalTool.ShortDescription, externalTool.Description)
|
||||||
description := externalTool.ShortDescription
|
|
||||||
if description == "" {
|
|
||||||
description = externalTool.Description
|
|
||||||
}
|
|
||||||
|
|
||||||
// 转换schema中的类型为OpenAI标准类型
|
// 转换schema中的类型为OpenAI标准类型
|
||||||
convertedSchema := a.convertSchemaTypes(externalTool.InputSchema)
|
convertedSchema := a.convertSchemaTypes(externalTool.InputSchema)
|
||||||
@@ -1063,6 +1057,19 @@ func (a *Agent) getAvailableTools(roleTools []string) []Tool {
|
|||||||
return tools
|
return tools
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func (a *Agent) pickToolDescription(shortDesc, fullDesc string) string {
|
||||||
|
a.mu.RLock()
|
||||||
|
mode := strings.TrimSpace(strings.ToLower(a.toolDescriptionMode))
|
||||||
|
a.mu.RUnlock()
|
||||||
|
if mode == "full" {
|
||||||
|
return fullDesc
|
||||||
|
}
|
||||||
|
if shortDesc != "" {
|
||||||
|
return shortDesc
|
||||||
|
}
|
||||||
|
return fullDesc
|
||||||
|
}
|
||||||
|
|
||||||
// convertSchemaTypes 递归转换schema中的类型为OpenAI标准类型
|
// convertSchemaTypes 递归转换schema中的类型为OpenAI标准类型
|
||||||
func (a *Agent) convertSchemaTypes(schema map[string]interface{}) map[string]interface{} {
|
func (a *Agent) convertSchemaTypes(schema map[string]interface{}) map[string]interface{} {
|
||||||
if schema == nil {
|
if schema == nil {
|
||||||
@@ -1665,6 +1672,18 @@ func (a *Agent) UpdateMaxIterations(maxIterations int) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// UpdateToolDescriptionMode 更新工具描述模式(short/full)
|
||||||
|
func (a *Agent) UpdateToolDescriptionMode(mode string) {
|
||||||
|
a.mu.Lock()
|
||||||
|
defer a.mu.Unlock()
|
||||||
|
mode = strings.TrimSpace(strings.ToLower(mode))
|
||||||
|
if mode != "full" {
|
||||||
|
mode = "short"
|
||||||
|
}
|
||||||
|
a.toolDescriptionMode = mode
|
||||||
|
a.logger.Info("Agent工具描述模式已更新", zap.String("tool_description_mode", mode))
|
||||||
|
}
|
||||||
|
|
||||||
// formatToolError 格式化工具错误信息,提供更友好的错误描述
|
// formatToolError 格式化工具错误信息,提供更友好的错误描述
|
||||||
func (a *Agent) formatToolError(toolName string, args map[string]interface{}, err error) string {
|
func (a *Agent) formatToolError(toolName string, args map[string]interface{}, err error) string {
|
||||||
errorMsg := fmt.Sprintf(`工具执行失败
|
errorMsg := fmt.Sprintf(`工具执行失败
|
||||||
|
|||||||
@@ -133,6 +133,7 @@ func New(cfg *config.Config, log *logger.Logger) (*App, error) {
|
|||||||
maxIterations = 30 // 默认值
|
maxIterations = 30 // 默认值
|
||||||
}
|
}
|
||||||
agent := agent.NewAgent(&cfg.OpenAI, &cfg.Agent, mcpServer, externalMCPMgr, log.Logger, maxIterations)
|
agent := agent.NewAgent(&cfg.OpenAI, &cfg.Agent, mcpServer, externalMCPMgr, log.Logger, maxIterations)
|
||||||
|
agent.UpdateToolDescriptionMode(cfg.Security.ToolDescriptionMode)
|
||||||
|
|
||||||
// 设置结果存储到Agent
|
// 设置结果存储到Agent
|
||||||
agent.SetResultStorage(resultStorage)
|
agent.SetResultStorage(resultStorage)
|
||||||
|
|||||||
@@ -90,6 +90,7 @@ type AttackChainUpdater interface {
|
|||||||
type AgentUpdater interface {
|
type AgentUpdater interface {
|
||||||
UpdateConfig(cfg *config.OpenAIConfig)
|
UpdateConfig(cfg *config.OpenAIConfig)
|
||||||
UpdateMaxIterations(maxIterations int)
|
UpdateMaxIterations(maxIterations int)
|
||||||
|
UpdateToolDescriptionMode(mode string)
|
||||||
}
|
}
|
||||||
|
|
||||||
// NewConfigHandler 创建新的配置处理器
|
// NewConfigHandler 创建新的配置处理器
|
||||||
@@ -232,13 +233,7 @@ func (h *ConfigHandler) GetConfig(c *gin.Context) {
|
|||||||
if configToolMap[mcpTool.Name] {
|
if configToolMap[mcpTool.Name] {
|
||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
description := mcpTool.ShortDescription
|
description := h.pickToolDescription(mcpTool.ShortDescription, mcpTool.Description)
|
||||||
if description == "" {
|
|
||||||
description = mcpTool.Description
|
|
||||||
}
|
|
||||||
if len(description) > 10000 {
|
|
||||||
description = description[:10000] + "..."
|
|
||||||
}
|
|
||||||
tools = append(tools, ToolConfigInfo{
|
tools = append(tools, ToolConfigInfo{
|
||||||
Name: mcpTool.Name,
|
Name: mcpTool.Name,
|
||||||
Description: description,
|
Description: description,
|
||||||
@@ -430,13 +425,7 @@ func (h *ConfigHandler) GetTools(c *gin.Context) {
|
|||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
|
|
||||||
description := mcpTool.ShortDescription
|
description := h.pickToolDescription(mcpTool.ShortDescription, mcpTool.Description)
|
||||||
if description == "" {
|
|
||||||
description = mcpTool.Description
|
|
||||||
}
|
|
||||||
if len(description) > 10000 {
|
|
||||||
description = description[:10000] + "..."
|
|
||||||
}
|
|
||||||
|
|
||||||
toolInfo := ToolConfigInfo{
|
toolInfo := ToolConfigInfo{
|
||||||
Name: mcpTool.Name,
|
Name: mcpTool.Name,
|
||||||
@@ -1061,6 +1050,7 @@ func (h *ConfigHandler) ApplyConfig(c *gin.Context) {
|
|||||||
if h.agent != nil {
|
if h.agent != nil {
|
||||||
h.agent.UpdateConfig(&h.config.OpenAI)
|
h.agent.UpdateConfig(&h.config.OpenAI)
|
||||||
h.agent.UpdateMaxIterations(h.config.Agent.MaxIterations)
|
h.agent.UpdateMaxIterations(h.config.Agent.MaxIterations)
|
||||||
|
h.agent.UpdateToolDescriptionMode(h.config.Security.ToolDescriptionMode)
|
||||||
h.logger.Info("Agent配置已更新")
|
h.logger.Info("Agent配置已更新")
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user