diff --git a/internal/multiagent/eino_summarize.go b/internal/multiagent/eino_summarize.go index dc80c868..4c40e906 100644 --- a/internal/multiagent/eino_summarize.go +++ b/internal/multiagent/eino_summarize.go @@ -57,19 +57,30 @@ func newEinoSummarizationMiddleware( if modelName == "" { modelName = "gpt-4o" } + tokenCounter := einoSummarizationTokenCounter(modelName) + recentTrailMax := trigger / 4 + if recentTrailMax < 2048 { + recentTrailMax = 2048 + } + if recentTrailMax > trigger/2 { + recentTrailMax = trigger / 2 + } mw, err := summarization.New(ctx, &summarization.Config{ Model: summaryModel, Trigger: &summarization.TriggerCondition{ ContextTokens: trigger, }, - TokenCounter: einoSummarizationTokenCounter(modelName), + TokenCounter: tokenCounter, UserInstruction: einoSummarizeUserInstruction, EmitInternalEvents: false, PreserveUserMessages: &summarization.PreserveUserMessages{ Enabled: true, MaxTokens: preserveMax, }, + Finalize: func(ctx context.Context, originalMessages []adk.Message, summary adk.Message) ([]adk.Message, error) { + return summarizeFinalizeWithRecentAssistantToolTrail(ctx, originalMessages, summary, tokenCounter, recentTrailMax) + }, Callback: func(ctx context.Context, before, after adk.ChatModelAgentState) error { if logger == nil { return nil @@ -89,6 +100,108 @@ func newEinoSummarizationMiddleware( return mw, nil } +// summarizeFinalizeWithRecentAssistantToolTrail 在摘要消息后保留最近 assistant/tool 轨迹,避免压缩后执行链断裂。 +func summarizeFinalizeWithRecentAssistantToolTrail( + ctx context.Context, + originalMessages []adk.Message, + summary adk.Message, + tokenCounter summarization.TokenCounterFunc, + recentTrailTokenBudget int, +) ([]adk.Message, error) { + systemMsgs := make([]adk.Message, 0, len(originalMessages)) + nonSystem := make([]adk.Message, 0, len(originalMessages)) + for _, msg := range originalMessages { + if msg == nil { + continue + } + if msg.Role == schema.System { + systemMsgs = append(systemMsgs, msg) + continue + } + nonSystem = append(nonSystem, msg) + } + + if recentTrailTokenBudget <= 0 || len(nonSystem) == 0 { + out := make([]adk.Message, 0, len(systemMsgs)+1) + out = append(out, systemMsgs...) + out = append(out, summary) + return out, nil + } + + selectedReverse := make([]adk.Message, 0, 8) + seen := make(map[adk.Message]struct{}) + totalTokens := 0 + assistantToolKept := 0 + const minAssistantToolTrail = 4 + + tryKeep := func(msg adk.Message) (bool, error) { + if msg == nil { + return false, nil + } + if _, ok := seen[msg]; ok { + return false, nil + } + n, err := tokenCounter(ctx, &summarization.TokenCounterInput{Messages: []adk.Message{msg}}) + if err != nil { + return false, err + } + if n <= 0 { + n = 1 + } + if totalTokens+n > recentTrailTokenBudget { + return false, nil + } + totalTokens += n + selectedReverse = append(selectedReverse, msg) + seen[msg] = struct{}{} + return true, nil + } + + // 优先保留最近 assistant/tool,确保执行轨迹可续跑。 + for i := len(nonSystem) - 1; i >= 0; i-- { + msg := nonSystem[i] + if msg.Role != schema.Assistant && msg.Role != schema.Tool { + continue + } + ok, err := tryKeep(msg) + if err != nil { + return nil, err + } + if ok { + assistantToolKept++ + } + if assistantToolKept >= minAssistantToolTrail { + break + } + } + + // 在预算内回填更多最近消息,保持短链路上下文。 + for i := len(nonSystem) - 1; i >= 0; i-- { + _, exists := seen[nonSystem[i]] + if exists { + continue + } + ok, err := tryKeep(nonSystem[i]) + if err != nil { + return nil, err + } + if !ok { + break + } + } + + selected := make([]adk.Message, 0, len(selectedReverse)) + for i := len(selectedReverse) - 1; i >= 0; i-- { + selected = append(selected, selectedReverse[i]) + } + + out := make([]adk.Message, 0, len(systemMsgs)+1+len(selected)) + out = append(out, systemMsgs...) + out = append(out, summary) + out = append(out, selected...) + return out, nil +} + func einoSummarizationTokenCounter(openAIModel string) summarization.TokenCounterFunc { tc := agent.NewTikTokenCounter() return func(ctx context.Context, input *summarization.TokenCounterInput) (int, error) {