CyberStrikeAI/internal/multiagent/eino_orchestration.go

package multiagent

import (
	"context"
	"fmt"
	"strings"

	"cyberstrike-ai/internal/agent"
	"cyberstrike-ai/internal/config"

	"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
	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
}

// 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,
	}
	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),
	})
	if err != nil {
		return nil, fmt.Errorf("plan_execute replanner: %w", err)
	}

	// 组装 executor handler 栈，顺序与 Deep/Supervisor 主代理一致（outermost first）。
	var execHandlers []adk.ChatModelAgentMiddleware
	// 1. patchtoolcalls, reduction, toolsearch, plantask（来自 prependEinoMiddlewares）
	if len(a.ExecPreMiddlewares) > 0 {
		execHandlers = append(execHandlers, a.ExecPreMiddlewares...)
	}
	// 2. filesystem 中间件（可选）
	if a.FilesystemMiddleware != nil {
		execHandlers = append(execHandlers, a.FilesystemMiddleware)
	}
	// 3. skill 中间件（可选）
	if a.SkillMiddleware != nil {
		execHandlers = append(execHandlers, a.SkillMiddleware)
	}
	// 4. summarization（最后，与 Deep/Supervisor 一致）
	if a.AppCfg != nil {
		sumMw, sumErr := newEinoSummarizationMiddleware(ctx, a.ExecModel, a.AppCfg, a.MwCfg, a.ConversationID, a.Logger)
		if sumErr != nil {
			return nil, fmt.Errorf("plan_execute executor summarization: %w", sumErr)
		}
		execHandlers = append(execHandlers, sumMw)
	}
	// 5. 孤儿 tool 消息兜底：必须挂在所有改写历史中间件（summarization/reduction/skill）之后、
	//    telemetry 之前，保证送入 ChatModel 的消息序列 tool_call ↔ tool_result 配对完整。
	execHandlers = append(execHandlers, newOrphanToolPrunerMiddleware(a.Logger, "plan_execute_executor"))
	if teleMw := newEinoModelInputTelemetryMiddleware(a.Logger, a.ModelName, a.ConversationID, "plan_execute_executor"); teleMw != nil {
		execHandlers = append(execHandlers, teleMw)
	}
	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,
	})
}

// 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, 1+len(userInput))
		if oi != "" {
			msgs = append(msgs, schema.SystemMessage(oi))
		}
		msgs = append(msgs, userInput...)
		if rewritten, rerr := applyBeforeModelRewriteHandlers(ctx, msgs, rewriteHandlers); rerr == nil && len(rewritten) > 0 {
			msgs = rewritten
		}
		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
		}
		if oi != "" {
			userMsgs = append([]adk.Message{schema.SystemMessage(oi)}, userMsgs...)
		}
		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 oi != "" {
			msgs = append([]adk.Message{schema.SystemMessage(oi)}, msgs...)
		}
		if rewritten, rerr := applyBeforeModelRewriteHandlers(ctx, msgs, rewriteHandlers); rerr == nil && len(rewritten) > 0 {
			msgs = rewritten
		}
		logPlanExecuteModelInputEstimate(logger, modelName, conversationID, "plan_execute_replanner", msgs)
		return msgs, nil
	}
}

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"
}