mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-08-28 21:50:43 +02:00
Add files via upload
This commit is contained in:
公明
@@ -0,0 +1,411 @@
|
||||
package multiagent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"cyberstrike-ai/internal/agent"
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/database"
|
||||
|
||||
"github.com/cloudwego/eino-ext/components/model/openai"
|
||||
"github.com/cloudwego/eino/adk"
|
||||
"github.com/cloudwego/eino/adk/prebuilt/planexecute"
|
||||
"github.com/cloudwego/eino/components/model"
|
||||
"github.com/cloudwego/eino/schema"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// PlanExecuteRootArgs 构建 Eino adk/prebuilt/planexecute 根 Agent 所需参数。
|
||||
type PlanExecuteRootArgs struct {
|
||||
MainToolCallingModel *openai.ChatModel
|
||||
ExecModel *openai.ChatModel
|
||||
OrchInstruction string
|
||||
ToolsCfg adk.ToolsConfig
|
||||
ExecMaxIter int
|
||||
LoopMaxIter int
|
||||
// AppCfg / Logger 非空时为 Executor 挂载与 Deep/Supervisor 一致的 Eino summarization 中间件。
|
||||
AppCfg *config.Config
|
||||
MwCfg *config.MultiAgentEinoMiddlewareConfig
|
||||
// ConversationID is used for transcript/isolation paths in middleware.
|
||||
ConversationID string
|
||||
DB *database.DB
|
||||
ProjectID string
|
||||
Logger *zap.Logger
|
||||
// ModelName is used for model input token estimation logs.
|
||||
ModelName string
|
||||
// ExecPreMiddlewares 是由 prependEinoMiddlewares 构建的前置中间件（patchtoolcalls, reduction, toolsearch, plantask），
|
||||
// 与 Deep/Supervisor 主代理的 mainOrchestratorPre 一致。
|
||||
ExecPreMiddlewares []adk.ChatModelAgentMiddleware
|
||||
// SkillMiddleware 是 Eino 官方 skill 渐进式披露中间件（可选）。
|
||||
SkillMiddleware adk.ChatModelAgentMiddleware
|
||||
// FilesystemMiddleware 是 Eino filesystem 中间件，当 eino_skills.filesystem_tools 启用时提供本机文件读写与 Shell 能力（可选）。
|
||||
FilesystemMiddleware adk.ChatModelAgentMiddleware
|
||||
// PlannerReplannerRewriteHandlers applies BeforeModelRewriteState pipeline for planner/replanner input.
|
||||
PlannerReplannerRewriteHandlers []adk.ChatModelAgentMiddleware
|
||||
// ModelFacingTrace 可选：由 Executor Handlers 链末尾写入，供 last_react 与 summarization 后上下文对齐。
|
||||
ModelFacingTrace *modelFacingTraceHolder
|
||||
}
|
||||
|
||||
// NewPlanExecuteRoot 返回 plan → execute → replan 预置编排根节点（与 Deep / Supervisor 并列）。
|
||||
func NewPlanExecuteRoot(ctx context.Context, a *PlanExecuteRootArgs) (adk.ResumableAgent, error) {
|
||||
if a == nil {
|
||||
return nil, fmt.Errorf("plan_execute: args 为空")
|
||||
}
|
||||
if a.MainToolCallingModel == nil || a.ExecModel == nil {
|
||||
return nil, fmt.Errorf("plan_execute: 模型为空")
|
||||
}
|
||||
tcm, ok := interface{}(a.MainToolCallingModel).(model.ToolCallingChatModel)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("plan_execute: 主模型需实现 ToolCallingChatModel")
|
||||
}
|
||||
plannerCfg := &planexecute.PlannerConfig{
|
||||
ToolCallingChatModel: tcm,
|
||||
NewPlan: newLenientPlan,
|
||||
}
|
||||
if fn := planExecutePlannerGenInput(a.OrchInstruction, a.AppCfg, a.MwCfg, a.Logger, a.ModelName, a.ConversationID, a.PlannerReplannerRewriteHandlers); fn != nil {
|
||||
plannerCfg.GenInputFn = fn
|
||||
}
|
||||
planner, err := planexecute.NewPlanner(ctx, plannerCfg)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("plan_execute planner: %w", err)
|
||||
}
|
||||
replanner, err := planexecute.NewReplanner(ctx, &planexecute.ReplannerConfig{
|
||||
ChatModel: tcm,
|
||||
GenInputFn: planExecuteReplannerGenInput(a.OrchInstruction, a.AppCfg, a.MwCfg, a.Logger, a.ModelName, a.ConversationID, a.PlannerReplannerRewriteHandlers),
|
||||
NewPlan: newLenientPlan,
|
||||
})
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("plan_execute replanner: %w", err)
|
||||
}
|
||||
|
||||
execHandlers, err := buildPlanExecuteExecutorHandlers(ctx, a)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
executor, err := newPlanExecuteExecutor(ctx, &planexecute.ExecutorConfig{
|
||||
Model: a.ExecModel,
|
||||
ToolsConfig: a.ToolsCfg,
|
||||
MaxIterations: a.ExecMaxIter,
|
||||
GenInputFn: planExecuteExecutorGenInput(a.OrchInstruction, a.AppCfg, a.MwCfg, a.Logger, a.ModelName, a.ConversationID),
|
||||
}, execHandlers)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("plan_execute executor: %w", err)
|
||||
}
|
||||
loopMax := a.LoopMaxIter
|
||||
if loopMax <= 0 {
|
||||
loopMax = 10
|
||||
}
|
||||
return planexecute.New(ctx, &planexecute.Config{
|
||||
Planner: planner,
|
||||
Executor: executor,
|
||||
Replanner: replanner,
|
||||
MaxIterations: loopMax,
|
||||
})
|
||||
}
|
||||
|
||||
// buildPlanExecuteExecutorHandlers 组装 Executor 中间件栈（outermost first），与 Deep/Supervisor 主代理对齐：
|
||||
// ExecPreMiddlewares（patch / reduction / toolsearch / plantask）→ filesystem → skill → summarization tail。
|
||||
func buildPlanExecuteExecutorHandlers(ctx context.Context, a *PlanExecuteRootArgs) ([]adk.ChatModelAgentMiddleware, error) {
|
||||
if a == nil {
|
||||
return nil, fmt.Errorf("plan_execute: args 为空")
|
||||
}
|
||||
var execHandlers []adk.ChatModelAgentMiddleware
|
||||
if len(a.ExecPreMiddlewares) > 0 {
|
||||
execHandlers = append(execHandlers, a.ExecPreMiddlewares...)
|
||||
}
|
||||
if a.FilesystemMiddleware != nil {
|
||||
execHandlers = append(execHandlers, a.FilesystemMiddleware)
|
||||
}
|
||||
if a.SkillMiddleware != nil {
|
||||
execHandlers = append(execHandlers, a.SkillMiddleware)
|
||||
}
|
||||
if a.AppCfg != nil {
|
||||
sumMw, sumErr := newEinoSummarizationMiddleware(ctx, a.ExecModel, a.AppCfg, a.MwCfg, a.ConversationID, a.DB, a.ProjectID, a.Logger)
|
||||
if sumErr != nil {
|
||||
return nil, fmt.Errorf("plan_execute executor summarization: %w", sumErr)
|
||||
}
|
||||
execHandlers = appendEinoChatModelTailMiddlewares(execHandlers, einoChatModelTailConfig{
|
||||
logger: a.Logger,
|
||||
phase: "plan_execute_executor",
|
||||
summarization: sumMw,
|
||||
modelName: a.ModelName,
|
||||
maxTotalTokens: a.AppCfg.OpenAI.MaxTotalTokens,
|
||||
toolMaxBytes: toolMaxBytesFromMW(a.MwCfg),
|
||||
conversationID: a.ConversationID,
|
||||
trace: a.ModelFacingTrace,
|
||||
})
|
||||
}
|
||||
return execHandlers, nil
|
||||
}
|
||||
|
||||
// planExecutePlannerGenInput 将 orchestrator instruction 作为 SystemMessage 注入 planner 输入。
|
||||
// 返回 nil 时 Eino 使用内置默认 planner prompt。
|
||||
func planExecutePlannerGenInput(
|
||||
orchInstruction string,
|
||||
appCfg *config.Config,
|
||||
mwCfg *config.MultiAgentEinoMiddlewareConfig,
|
||||
logger *zap.Logger,
|
||||
modelName string,
|
||||
conversationID string,
|
||||
rewriteHandlers []adk.ChatModelAgentMiddleware,
|
||||
) planexecute.GenPlannerModelInputFn {
|
||||
oi := strings.TrimSpace(orchInstruction)
|
||||
if oi == "" && appCfg == nil {
|
||||
return nil
|
||||
}
|
||||
return func(ctx context.Context, userInput []adk.Message) ([]adk.Message, error) {
|
||||
userInput = capPlanExecuteUserInputMessages(userInput, appCfg, mwCfg)
|
||||
msgs := make([]adk.Message, 0, len(userInput))
|
||||
msgs = append(msgs, userInput...)
|
||||
if rewritten, rerr := applyBeforeModelRewriteHandlers(ctx, msgs, rewriteHandlers); rerr == nil && len(rewritten) > 0 {
|
||||
msgs = rewritten
|
||||
}
|
||||
msgs = normalizeSingleLeadingSystemMessage(msgs, oi)
|
||||
logPlanExecuteModelInputEstimate(logger, modelName, conversationID, "plan_execute_planner", msgs)
|
||||
return msgs, nil
|
||||
}
|
||||
}
|
||||
|
||||
func planExecuteExecutorGenInput(
|
||||
orchInstruction string,
|
||||
appCfg *config.Config,
|
||||
mwCfg *config.MultiAgentEinoMiddlewareConfig,
|
||||
logger *zap.Logger,
|
||||
modelName string,
|
||||
conversationID string,
|
||||
) planexecute.GenModelInputFn {
|
||||
oi := strings.TrimSpace(orchInstruction)
|
||||
return func(ctx context.Context, in *planexecute.ExecutionContext) ([]adk.Message, error) {
|
||||
planContent, err := in.Plan.MarshalJSON()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
userMsgs, err := planexecute.ExecutorPrompt.Format(ctx, map[string]any{
|
||||
"input": planExecuteFormatInput(capPlanExecuteUserInputMessages(in.UserInput, appCfg, mwCfg)),
|
||||
"plan": string(planContent),
|
||||
"executed_steps": planExecuteFormatExecutedSteps(in.ExecutedSteps, appCfg, mwCfg),
|
||||
"step": in.Plan.FirstStep(),
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
userMsgs = normalizeSingleLeadingSystemMessage(userMsgs, oi)
|
||||
logPlanExecuteModelInputEstimate(logger, modelName, conversationID, "plan_execute_executor_gen_input", userMsgs)
|
||||
return userMsgs, nil
|
||||
}
|
||||
}
|
||||
|
||||
func planExecuteFormatInput(input []adk.Message) string {
|
||||
var sb strings.Builder
|
||||
for _, msg := range input {
|
||||
sb.WriteString(msg.Content)
|
||||
sb.WriteString("\n")
|
||||
}
|
||||
return sb.String()
|
||||
}
|
||||
|
||||
func planExecuteFormatExecutedSteps(results []planexecute.ExecutedStep, appCfg *config.Config, mwCfg *config.MultiAgentEinoMiddlewareConfig) string {
|
||||
capped := capPlanExecuteExecutedStepsWithConfig(results, mwCfg)
|
||||
return renderPlanExecuteStepsByBudget(capped, appCfg, mwCfg)
|
||||
}
|
||||
|
||||
// planExecuteReplannerGenInput 与 Eino 默认 Replanner 输入一致，但 executed_steps 经 cap 后再写入 prompt，
|
||||
// 且在 orchInstruction 非空时 prepend SystemMessage 使 replanner 也能接收全局指令。
|
||||
func planExecuteReplannerGenInput(
|
||||
orchInstruction string,
|
||||
appCfg *config.Config,
|
||||
mwCfg *config.MultiAgentEinoMiddlewareConfig,
|
||||
logger *zap.Logger,
|
||||
modelName string,
|
||||
conversationID string,
|
||||
rewriteHandlers []adk.ChatModelAgentMiddleware,
|
||||
) planexecute.GenModelInputFn {
|
||||
oi := strings.TrimSpace(orchInstruction)
|
||||
return func(ctx context.Context, in *planexecute.ExecutionContext) ([]adk.Message, error) {
|
||||
planContent, err := in.Plan.MarshalJSON()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
msgs, err := planexecute.ReplannerPrompt.Format(ctx, map[string]any{
|
||||
"plan": string(planContent),
|
||||
"input": planExecuteFormatInput(capPlanExecuteUserInputMessages(in.UserInput, appCfg, mwCfg)),
|
||||
"executed_steps": planExecuteFormatExecutedSteps(in.ExecutedSteps, appCfg, mwCfg),
|
||||
"plan_tool": planexecute.PlanToolInfo.Name,
|
||||
"respond_tool": planexecute.RespondToolInfo.Name,
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if rewritten, rerr := applyBeforeModelRewriteHandlers(ctx, msgs, rewriteHandlers); rerr == nil && len(rewritten) > 0 {
|
||||
msgs = rewritten
|
||||
}
|
||||
msgs = normalizeSingleLeadingSystemMessage(msgs, oi)
|
||||
logPlanExecuteModelInputEstimate(logger, modelName, conversationID, "plan_execute_replanner", msgs)
|
||||
return msgs, nil
|
||||
}
|
||||
}
|
||||
|
||||
// normalizeSingleLeadingSystemMessage enforces a provider-friendly message shape:
|
||||
// exactly one system message at index 0 (when any system context exists).
|
||||
// For strict OpenAI-compatible backends (e.g. qwen/vllm templates), this avoids
|
||||
// "System message must be at the beginning" caused by multiple/disordered system messages.
|
||||
func normalizeSingleLeadingSystemMessage(msgs []adk.Message, extraSystem string) []adk.Message {
|
||||
extraSystem = strings.TrimSpace(extraSystem)
|
||||
if len(msgs) == 0 {
|
||||
if extraSystem == "" {
|
||||
return msgs
|
||||
}
|
||||
return []adk.Message{schema.SystemMessage(extraSystem)}
|
||||
}
|
||||
|
||||
systemParts := make([]string, 0, 2)
|
||||
if extraSystem != "" {
|
||||
systemParts = append(systemParts, extraSystem)
|
||||
}
|
||||
nonSystem := make([]adk.Message, 0, len(msgs))
|
||||
for _, msg := range msgs {
|
||||
if msg == nil {
|
||||
continue
|
||||
}
|
||||
if msg.Role == schema.System {
|
||||
if s := strings.TrimSpace(msg.Content); s != "" {
|
||||
systemParts = append(systemParts, s)
|
||||
}
|
||||
continue
|
||||
}
|
||||
nonSystem = append(nonSystem, msg)
|
||||
}
|
||||
if len(systemParts) == 0 {
|
||||
return nonSystem
|
||||
}
|
||||
out := make([]adk.Message, 0, len(nonSystem)+1)
|
||||
out = append(out, schema.SystemMessage(strings.Join(systemParts, "\n\n")))
|
||||
out = append(out, nonSystem...)
|
||||
return out
|
||||
}
|
||||
|
||||
func capPlanExecuteUserInputMessages(input []adk.Message, appCfg *config.Config, mwCfg *config.MultiAgentEinoMiddlewareConfig) []adk.Message {
|
||||
if len(input) == 0 {
|
||||
return input
|
||||
}
|
||||
maxTotal := 120000
|
||||
modelName := "gpt-4o"
|
||||
if appCfg != nil {
|
||||
if appCfg.OpenAI.MaxTotalTokens > 0 {
|
||||
maxTotal = appCfg.OpenAI.MaxTotalTokens
|
||||
}
|
||||
if m := strings.TrimSpace(appCfg.OpenAI.Model); m != "" {
|
||||
modelName = m
|
||||
}
|
||||
}
|
||||
// Reserve most tokens for planner/replanner prompt and tool schema.
|
||||
ratio := 0.35
|
||||
if mwCfg != nil {
|
||||
ratio = mwCfg.PlanExecuteUserInputBudgetRatioEffective()
|
||||
}
|
||||
budget := int(float64(maxTotal) * ratio)
|
||||
if budget < 4096 {
|
||||
budget = 4096
|
||||
}
|
||||
tc := agent.NewTikTokenCounter()
|
||||
out := make([]adk.Message, 0, len(input))
|
||||
used := 0
|
||||
for i := len(input) - 1; i >= 0; i-- {
|
||||
msg := input[i]
|
||||
if msg == nil {
|
||||
continue
|
||||
}
|
||||
n, err := tc.Count(modelName, string(msg.Role)+"\n"+msg.Content)
|
||||
if err != nil {
|
||||
n = (len(msg.Content) + 3) / 4
|
||||
}
|
||||
if n <= 0 {
|
||||
n = 1
|
||||
}
|
||||
if used+n > budget {
|
||||
break
|
||||
}
|
||||
used += n
|
||||
out = append(out, msg)
|
||||
}
|
||||
for i, j := 0, len(out)-1; i < j; i, j = i+1, j-1 {
|
||||
out[i], out[j] = out[j], out[i]
|
||||
}
|
||||
if len(out) == 0 {
|
||||
// Keep the latest user message at least.
|
||||
return []adk.Message{input[len(input)-1]}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func renderPlanExecuteStepsByBudget(steps []planexecute.ExecutedStep, appCfg *config.Config, mwCfg *config.MultiAgentEinoMiddlewareConfig) string {
|
||||
if len(steps) == 0 {
|
||||
return ""
|
||||
}
|
||||
maxTotal := 120000
|
||||
modelName := "gpt-4o"
|
||||
if appCfg != nil {
|
||||
if appCfg.OpenAI.MaxTotalTokens > 0 {
|
||||
maxTotal = appCfg.OpenAI.MaxTotalTokens
|
||||
}
|
||||
if m := strings.TrimSpace(appCfg.OpenAI.Model); m != "" {
|
||||
modelName = m
|
||||
}
|
||||
}
|
||||
ratio := 0.2
|
||||
if mwCfg != nil {
|
||||
ratio = mwCfg.PlanExecuteExecutedStepsBudgetRatioEffective()
|
||||
}
|
||||
budget := int(float64(maxTotal) * ratio)
|
||||
if budget < 3072 {
|
||||
budget = 3072
|
||||
}
|
||||
tc := agent.NewTikTokenCounter()
|
||||
var kept []string
|
||||
used := 0
|
||||
skipped := 0
|
||||
for i := len(steps) - 1; i >= 0; i-- {
|
||||
block := fmt.Sprintf("Step: %s\nResult: %s\n\n", steps[i].Step, steps[i].Result)
|
||||
n, err := tc.Count(modelName, block)
|
||||
if err != nil {
|
||||
n = (len(block) + 3) / 4
|
||||
}
|
||||
if n <= 0 {
|
||||
n = 1
|
||||
}
|
||||
if used+n > budget {
|
||||
skipped = i + 1
|
||||
break
|
||||
}
|
||||
used += n
|
||||
kept = append(kept, block)
|
||||
}
|
||||
var sb strings.Builder
|
||||
if skipped > 0 {
|
||||
sb.WriteString(fmt.Sprintf("Earlier executed steps omitted due to context budget: %d steps.\n\n", skipped))
|
||||
}
|
||||
for i := len(kept) - 1; i >= 0; i-- {
|
||||
sb.WriteString(kept[i])
|
||||
}
|
||||
return sb.String()
|
||||
}
|
||||
|
||||
// planExecuteStreamsMainAssistant 将规划/执行/重规划各阶段助手流式输出映射到主对话区。
|
||||
func planExecuteStreamsMainAssistant(agent string) bool {
|
||||
if agent == "" {
|
||||
return true
|
||||
}
|
||||
switch agent {
|
||||
case "planner", "executor", "replanner", "execute_replan", "plan_execute_replan":
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func planExecuteEinoRoleTag(agent string) string {
|
||||
_ = agent
|
||||
return "orchestrator"
|
||||
}
|
||||
Reference in New Issue
Block a user