CalvinBackup/CyberStrikeAI
1
0
Fork 0
mirror of https://github.com/Ed1s0nZ/CyberStrikeAI.git synced 2026-04-01 00:30:33 +02:00
Code Issues Packages Projects Releases Wiki Activity

Add files via upload

Browse Source
This commit is contained in:
公明
2025-12-17 23:46:47 +08:00
committed by GitHub
parent ff3c910173
commit 95fc939b5a
4 changed files with 44 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
internal/agent/agent.go
View File

@@ -98,9 +98,10 @@ func NewAgent(cfg *config.OpenAIConfig, agentCfg *config.AgentConfig, mcpServer
var memoryCompressor *MemoryCompressor
if cfg != nil {
mc, err := NewMemoryCompressor(MemoryCompressorConfig{
OpenAIConfig: cfg,
HTTPClient: httpClient,
Logger: logger,
MaxTotalTokens: cfg.MaxTotalTokens,
OpenAIConfig: cfg,
HTTPClient: httpClient,
Logger: logger,
})
if err != nil {
logger.Warn("初始化MemoryCompressor失败，将跳过上下文压缩", zap.Error(err))

6
internal/agent/memory_compressor.go
View File

@@ -17,7 +17,6 @@ import (
)
const (
DefaultMaxTotalTokens = 120_000
DefaultMinRecentMessage = 10
defaultChunkSize = 10
defaultMaxImages = 3
@@ -86,9 +85,8 @@ func NewMemoryCompressor(cfg MemoryCompressorConfig) (*MemoryCompressor, error)
cfg.Logger = zap.NewNop()
}
if cfg.MaxTotalTokens <= 0 {
cfg.MaxTotalTokens = DefaultMaxTotalTokens
}
// 如果没有显式配置 MaxTotalTokens，则后续逻辑会根据模型的最大上下文长度进行控制；
// 优先推荐在 config.yaml 的 openai.max_total_tokens 中统一配置。
if cfg.MinRecentMessage <= 0 {
cfg.MinRecentMessage = DefaultMinRecentMessage
}

14
internal/attackchain/builder.go
View File

@@ -51,9 +51,12 @@ func NewBuilder(db *database.DB, openAIConfig *config.OpenAIConfig, logger *zap.
}
httpClient := &http.Client{Timeout: 5 * time.Minute, Transport: transport}
maxTokens := 100000 // 默认100k tokens，可以根据模型调整
// 根据模型设置合理的默认值
if openAIConfig != nil {
// 优先使用配置文件中的统一 Token 上限（config.yaml -> openai.max_total_tokens）
maxTokens := 0
if openAIConfig != nil && openAIConfig.MaxTotalTokens > 0 {
maxTokens = openAIConfig.MaxTotalTokens
} else if openAIConfig != nil {
// 如果未显式配置 max_total_tokens，则根据模型设置一个合理的默认值
model := strings.ToLower(openAIConfig.Model)
if strings.Contains(model, "gpt-4") {
maxTokens = 128000 // gpt-4通常支持128k
@@ -61,7 +64,12 @@ func NewBuilder(db *database.DB, openAIConfig *config.OpenAIConfig, logger *zap.
maxTokens = 16000 // gpt-3.5-turbo通常支持16k
} else if strings.Contains(model, "deepseek") {
maxTokens = 131072 // deepseek-chat通常支持131k
} else {
maxTokens = 100000 // 兜底默认值
}
} else {
// 没有 OpenAI 配置时使用兜底值，避免为 0
maxTokens = 100000
}
return &Builder{

52
internal/config/config.go
View File

@@ -12,15 +12,15 @@ import (
)
type Config struct {
Server ServerConfig `yaml:"server"`
Log LogConfig `yaml:"log"`
MCP MCPConfig `yaml:"mcp"`
OpenAI OpenAIConfig `yaml:"openai"`
Agent AgentConfig `yaml:"agent"`
Security SecurityConfig `yaml:"security"`
Database DatabaseConfig `yaml:"database"`
Auth AuthConfig `yaml:"auth"`
ExternalMCP ExternalMCPConfig `yaml:"external_mcp,omitempty"`
Server ServerConfig `yaml:"server"`
Log LogConfig `yaml:"log"`
MCP MCPConfig `yaml:"mcp"`
OpenAI OpenAIConfig `yaml:"openai"`
Agent AgentConfig `yaml:"agent"`
Security SecurityConfig `yaml:"security"`
Database DatabaseConfig `yaml:"database"`
Auth AuthConfig `yaml:"auth"`
ExternalMCP ExternalMCPConfig `yaml:"external_mcp,omitempty"`
}
type ServerConfig struct {
@@ -40,9 +40,10 @@ type MCPConfig struct {
}
type OpenAIConfig struct {
APIKey string `yaml:"api_key" json:"api_key"`
BaseURL string `yaml:"base_url" json:"base_url"`
Model string `yaml:"model" json:"model"`
APIKey string `yaml:"api_key" json:"api_key"`
BaseURL string `yaml:"base_url" json:"base_url"`
Model string `yaml:"model" json:"model"`
MaxTotalTokens int `yaml:"max_total_tokens,omitempty" json:"max_total_tokens,omitempty"`
}
type SecurityConfig struct {
@@ -76,19 +77,19 @@ type ExternalMCPConfig struct {
// ExternalMCPServerConfig 外部MCP服务器配置
type ExternalMCPServerConfig struct {
// stdio模式配置
Command string `yaml:"command,omitempty" json:"command,omitempty"`
Args []string `yaml:"args,omitempty" json:"args,omitempty"`
Command string `yaml:"command,omitempty" json:"command,omitempty"`
Args []string `yaml:"args,omitempty" json:"args,omitempty"`
// HTTP模式配置
Transport string `yaml:"transport,omitempty" json:"transport,omitempty"` // "http" 或 "stdio"
URL string `yaml:"url,omitempty" json:"url,omitempty"`
Transport string `yaml:"transport,omitempty" json:"transport,omitempty"` // "http" 或 "stdio"
URL string `yaml:"url,omitempty" json:"url,omitempty"`
// 通用配置
Description string `yaml:"description,omitempty" json:"description,omitempty"`
Timeout int `yaml:"timeout,omitempty" json:"timeout,omitempty"` // 超时时间（秒）
ExternalMCPEnable bool `yaml:"external_mcp_enable,omitempty" json:"external_mcp_enable,omitempty"` // 是否启用外部MCP
ToolEnabled map[string]bool `yaml:"tool_enabled,omitempty" json:"tool_enabled,omitempty"` // 每个工具的启用状态（工具名称 -> 是否启用）
Description string `yaml:"description,omitempty" json:"description,omitempty"`
Timeout int `yaml:"timeout,omitempty" json:"timeout,omitempty"` // 超时时间（秒）
ExternalMCPEnable bool `yaml:"external_mcp_enable,omitempty" json:"external_mcp_enable,omitempty"` // 是否启用外部MCP
ToolEnabled map[string]bool `yaml:"tool_enabled,omitempty" json:"tool_enabled,omitempty"` // 每个工具的启用状态（工具名称 -> 是否启用）
// 向后兼容字段（已废弃，保留用于读取旧配置）
Enabled bool `yaml:"enabled,omitempty" json:"enabled,omitempty"` // 已废弃，使用 external_mcp_enable
Disabled bool `yaml:"disabled,omitempty" json:"disabled,omitempty"` // 已废弃，使用 external_mcp_enable
@@ -386,8 +387,9 @@ func Default() *Config {
Port: 8081,
},
OpenAI: OpenAIConfig{
BaseURL: "https://api.openai.com/v1",
Model: "gpt-4",
BaseURL: "https://api.openai.com/v1",
Model: "gpt-4",
MaxTotalTokens: 120000,
},
Agent: AgentConfig{
MaxIterations: 30, // 默认最大迭代次数