From 95fc939b5ac02311799b4843214d3eb3beac51d4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=85=AC=E6=98=8E?= <83812544+Ed1s0nZ@users.noreply.github.com> Date: Wed, 17 Dec 2025 23:46:47 +0800 Subject: [PATCH] Add files via upload --- internal/agent/agent.go | 7 ++-- internal/agent/memory_compressor.go | 6 ++-- internal/attackchain/builder.go | 14 ++++++-- internal/config/config.go | 52 +++++++++++++++-------------- 4 files changed, 44 insertions(+), 35 deletions(-) diff --git a/internal/agent/agent.go b/internal/agent/agent.go index 02ceb668..171bf82f 100644 --- a/internal/agent/agent.go +++ b/internal/agent/agent.go @@ -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)) diff --git a/internal/agent/memory_compressor.go b/internal/agent/memory_compressor.go index 3f1ac66f..edab2523 100644 --- a/internal/agent/memory_compressor.go +++ b/internal/agent/memory_compressor.go @@ -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 } diff --git a/internal/attackchain/builder.go b/internal/attackchain/builder.go index 394454dc..1e0a1392 100644 --- a/internal/attackchain/builder.go +++ b/internal/attackchain/builder.go @@ -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{ diff --git a/internal/config/config.go b/internal/config/config.go index 5dc2eacf..c8d61e4c 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -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, // 默认最大迭代次数