ARTEX — AI 自主渗透测试系统(Go 后端 + Next.js 前端,单二进制内嵌)。 Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
213 lines
6.8 KiB
Go
213 lines
6.8 KiB
Go
// Package agent wires real LLM-driven planner and work agents (on top of the
|
||
// agent-core SDK) to the dual SQLite graph. See docs/ARTEX-架构设计.md
|
||
// §4.3 (planner) and §4.4 (work agent).
|
||
//
|
||
// Provider configuration is read from the environment so the system runs with
|
||
// any Anthropic- or OpenAI-format endpoint. If no key is configured, FromEnv
|
||
// returns ok=false and the exploration engine stays idle (an LLM is required).
|
||
package agent
|
||
|
||
import (
|
||
"context"
|
||
"os"
|
||
"strings"
|
||
"time"
|
||
|
||
"github.com/Autumn-27/norma/agentcore"
|
||
"github.com/Autumn-27/norma/compaction"
|
||
"github.com/Autumn-27/norma/llm"
|
||
acperm "github.com/Autumn-27/norma/permission"
|
||
)
|
||
|
||
// Config describes the LLM backend resolved from the environment.
|
||
type Config struct {
|
||
Format llm.Format
|
||
BaseURL string
|
||
APIKey string
|
||
Model string
|
||
// Proxy routes all LLM requests through the given proxy URL (http/https/socks5).
|
||
// Empty falls back to the standard *_PROXY environment variables.
|
||
Proxy string
|
||
// RatePerSecond / RatePerMinute cap the shared request rate across ALL agents
|
||
// using the provider (0 = that window unlimited).
|
||
RatePerSecond float64
|
||
RatePerMinute float64
|
||
// ContextWindowK is the model's context window in K tokens (user-configured),
|
||
// used to size compaction thresholds. 0 = default; see CompactionWindow.
|
||
ContextWindowK int
|
||
// ReasoningEffort selects the thinking mode: "" = 默认(不发送思考参数);
|
||
// "off" = 关闭(thinking.type=disabled); "low"/"medium"/"high"/"max" = 开启并设强度.
|
||
// NewProvider derives the provider-specific request fields from it.
|
||
ReasoningEffort string
|
||
}
|
||
|
||
// compaction window resolution bounds (in K tokens). Below the floor the
|
||
// threshold math (window − summary reserve − buffer) would go non-positive and
|
||
// compaction would fire every turn; above the cap it would never fire.
|
||
const (
|
||
defaultWindowK = 200 // unset → assume a 200K window (Claude default)
|
||
minWindowK = 32 // floor so effectiveWindow stays comfortably positive
|
||
maxWindowK = 1000 // cap at 1M tokens (user request)
|
||
)
|
||
|
||
// CompactionWindow returns the model context window in TOKENS for compaction
|
||
// thresholds, resolved from the user-configured size (ContextWindowK). 0/unset →
|
||
// a 200K default; otherwise clamped to [32K, 1M] so compaction stays effective.
|
||
func (c Config) CompactionWindow() int {
|
||
k := c.ContextWindowK
|
||
if k <= 0 {
|
||
k = defaultWindowK
|
||
}
|
||
if k < minWindowK {
|
||
k = minWindowK
|
||
}
|
||
if k > maxWindowK {
|
||
k = maxWindowK
|
||
}
|
||
return k * 1000
|
||
}
|
||
|
||
// compactionConfig builds the agent-core compaction config for a context window
|
||
// in tokens. agentcore.NewSession wires the summarizer (same provider) when this
|
||
// is set on Options.Compaction.
|
||
func compactionConfig(windowTokens int) *compaction.Config {
|
||
if windowTokens <= 0 {
|
||
windowTokens = defaultWindowK * 1000
|
||
}
|
||
return &compaction.Config{ContextWindow: windowTokens}
|
||
}
|
||
|
||
// FromEnv resolves the LLM provider config:
|
||
//
|
||
// ARTEX_LLM_PROVIDER = anthropic|openai (default: inferred from keys)
|
||
// ARTEX_LLM_MODEL = model id (default: per provider)
|
||
// ARTEX_LLM_BASE_URL = endpoint (optional)
|
||
// ARTEX_LLM_PROXY = proxy URL (optional; http/https/socks5)
|
||
// ANTHROPIC_API_KEY / OPENAI_API_KEY = credentials
|
||
func FromEnv() (Config, bool) {
|
||
prov := os.Getenv("ARTEX_LLM_PROVIDER")
|
||
anthKey := os.Getenv("ANTHROPIC_API_KEY")
|
||
oaiKey := os.Getenv("OPENAI_API_KEY")
|
||
|
||
if prov == "" {
|
||
switch {
|
||
case anthKey != "":
|
||
prov = "anthropic"
|
||
case oaiKey != "":
|
||
prov = "openai"
|
||
default:
|
||
return Config{}, false
|
||
}
|
||
}
|
||
|
||
c := Config{
|
||
BaseURL: os.Getenv("ARTEX_LLM_BASE_URL"),
|
||
Model: os.Getenv("ARTEX_LLM_MODEL"),
|
||
Proxy: strings.TrimSpace(os.Getenv("ARTEX_LLM_PROXY")),
|
||
}
|
||
switch prov {
|
||
case "openai":
|
||
c.Format = llm.FormatOpenAI
|
||
c.APIKey = oaiKey
|
||
if c.Model == "" {
|
||
c.Model = "gpt-4o"
|
||
}
|
||
default:
|
||
c.Format = llm.FormatAnthropic
|
||
c.APIKey = anthKey
|
||
if c.Model == "" {
|
||
c.Model = "claude-opus-4-8"
|
||
}
|
||
}
|
||
if c.APIKey == "" {
|
||
return Config{}, false
|
||
}
|
||
return c, true
|
||
}
|
||
|
||
// ConfigFrom builds a Config from UI-provided strings (provider defaults to
|
||
// anthropic; model defaults per provider). Inputs are trimmed and the base URL
|
||
// is normalized to the API base the provider expects (the provider appends the
|
||
// endpoint path itself), so a full endpoint URL is tolerated.
|
||
func ConfigFrom(provider, model, baseURL, apiKey, proxy string) Config {
|
||
c := Config{
|
||
Model: strings.TrimSpace(model),
|
||
BaseURL: strings.TrimRight(strings.TrimSpace(baseURL), "/"),
|
||
APIKey: strings.TrimSpace(apiKey),
|
||
Proxy: strings.TrimSpace(proxy),
|
||
}
|
||
switch strings.TrimSpace(provider) {
|
||
case "openai":
|
||
c.Format = llm.FormatOpenAI
|
||
// provider appends "/chat/completions"; tolerate a full endpoint URL.
|
||
c.BaseURL = strings.TrimRight(strings.TrimSuffix(c.BaseURL, "/chat/completions"), "/")
|
||
if c.Model == "" {
|
||
c.Model = "gpt-4o"
|
||
}
|
||
default:
|
||
c.Format = llm.FormatAnthropic
|
||
// provider appends "/v1/messages".
|
||
c.BaseURL = strings.TrimRight(strings.TrimSuffix(c.BaseURL, "/v1/messages"), "/")
|
||
if c.Model == "" {
|
||
c.Model = "claude-opus-4-8"
|
||
}
|
||
}
|
||
return c
|
||
}
|
||
|
||
// Provider returns the short provider name ("anthropic"/"openai").
|
||
func (c Config) Provider() string {
|
||
if c.Format == llm.FormatOpenAI {
|
||
return "openai"
|
||
}
|
||
return "anthropic"
|
||
}
|
||
|
||
// NewProvider builds an llm.Provider from the config. When a rate is set, the
|
||
// limiter lives on the single provider instance — so planner + all workers +
|
||
// main agent (which share this provider) are bounded by one shared rate limit.
|
||
func (c Config) NewProvider() (llm.Provider, error) {
|
||
lc := llm.Config{
|
||
Format: c.Format,
|
||
BaseURL: c.BaseURL,
|
||
APIKey: c.APIKey,
|
||
Model: c.Model,
|
||
Proxy: c.Proxy,
|
||
}
|
||
// Derive the provider thinking params from the single UI selector:
|
||
// "" → omit both (provider default); "off" → disable; else → enable + effort.
|
||
switch c.ReasoningEffort {
|
||
case "":
|
||
// leave both empty → omitted from the request
|
||
case "off":
|
||
lc.ThinkingType = "disabled"
|
||
default:
|
||
lc.ThinkingType = "enabled"
|
||
lc.ReasoningEffort = c.ReasoningEffort
|
||
}
|
||
if c.RatePerSecond > 0 || c.RatePerMinute > 0 {
|
||
lc.RateLimit = &llm.RateLimit{PerSecond: c.RatePerSecond, PerMinute: c.RatePerMinute}
|
||
}
|
||
return llm.NewProvider(lc)
|
||
}
|
||
|
||
// TestConnection makes a minimal real completion to verify the provider/model/
|
||
// endpoint/key actually work. Returns the round-trip latency.
|
||
func TestConnection(ctx context.Context, c Config) (time.Duration, error) {
|
||
prov, err := c.NewProvider()
|
||
if err != nil {
|
||
return 0, err
|
||
}
|
||
ctx, cancel := context.WithTimeout(ctx, 30*time.Second)
|
||
defer cancel()
|
||
start := time.Now()
|
||
_, err = agentcore.Run(ctx, agentcore.Options{
|
||
Provider: prov,
|
||
SystemPrompt: []string{"你是连接测试。只回复 OK,不要别的。"},
|
||
PermissionMode: acperm.ModeBypass,
|
||
MaxTurns: 1,
|
||
MaxTokens: 32,
|
||
}, "ping")
|
||
return time.Since(start), err
|
||
}
|