mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-07-14 07:57:26 +02:00
Add files via upload
This commit is contained in:
公明
@@ -0,0 +1,328 @@
|
||||
package multiagent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
"unicode/utf8"
|
||||
|
||||
"github.com/cloudwego/eino/adk"
|
||||
"github.com/cloudwego/eino/adk/middlewares/summarization"
|
||||
"github.com/cloudwego/eino/schema"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
const (
|
||||
toolOutputTruncationMarker = "\n\n...[tool output truncated; full text persisted in reduction cache or summarization transcript]...\n\n"
|
||||
aggressiveToolTruncDivisor = 4
|
||||
)
|
||||
|
||||
// isEinoContextOverflowError reports API-side context window rejections.
|
||||
func isEinoContextOverflowError(err error) bool {
|
||||
if err == nil {
|
||||
return false
|
||||
}
|
||||
msg := strings.ToLower(strings.TrimSpace(err.Error()))
|
||||
if msg == "" {
|
||||
return false
|
||||
}
|
||||
markers := []string{
|
||||
"context length",
|
||||
"context_length",
|
||||
"maximum context",
|
||||
"max context",
|
||||
"context window",
|
||||
"context overflow",
|
||||
"too many tokens",
|
||||
"token limit",
|
||||
"tokens exceed",
|
||||
"exceeds the context",
|
||||
"input is too long",
|
||||
"prompt is too long",
|
||||
"request too large",
|
||||
}
|
||||
for _, m := range markers {
|
||||
if strings.Contains(msg, m) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func truncateBytesWithMarker(content string, maxBytes int, marker string) string {
|
||||
if maxBytes <= 0 || len(content) <= maxBytes {
|
||||
return content
|
||||
}
|
||||
if marker == "" {
|
||||
marker = toolOutputTruncationMarker
|
||||
}
|
||||
budget := maxBytes - len(marker)
|
||||
if budget <= 0 {
|
||||
if len(marker) > maxBytes {
|
||||
return marker[:maxBytes]
|
||||
}
|
||||
return marker
|
||||
}
|
||||
head := budget / 2
|
||||
tail := budget - head
|
||||
for head > 0 && !utf8.RuneStart(content[head]) {
|
||||
head--
|
||||
}
|
||||
tailStart := len(content) - tail
|
||||
for tailStart < len(content) && !utf8.RuneStart(content[tailStart]) {
|
||||
tailStart++
|
||||
}
|
||||
return content[:head] + marker + content[tailStart:]
|
||||
}
|
||||
|
||||
func cloneMessage(msg adk.Message) adk.Message {
|
||||
if msg == nil {
|
||||
return nil
|
||||
}
|
||||
cloned := *msg
|
||||
return &cloned
|
||||
}
|
||||
|
||||
func truncateMessageToolContent(msg adk.Message, maxBytes int, spillRef string) adk.Message {
|
||||
if msg == nil || maxBytes <= 0 {
|
||||
return msg
|
||||
}
|
||||
out := cloneMessage(msg)
|
||||
marker := toolOutputTruncationMarker
|
||||
if spillRef != "" {
|
||||
marker = fmt.Sprintf("\n\n...[tool output truncated; retrieve full text via: %s]...\n\n", spillRef)
|
||||
}
|
||||
switch out.Role {
|
||||
case schema.Tool:
|
||||
out.Content = truncateBytesWithMarker(out.Content, maxBytes, marker)
|
||||
case schema.Assistant:
|
||||
if out.ReasoningContent != "" {
|
||||
out.ReasoningContent = truncateBytesWithMarker(out.ReasoningContent, maxBytes, marker)
|
||||
}
|
||||
if out.Content != "" {
|
||||
out.Content = truncateBytesWithMarker(out.Content, maxBytes, marker)
|
||||
}
|
||||
case schema.User:
|
||||
if out.Content != "" {
|
||||
out.Content = truncateBytesWithMarker(out.Content, maxBytes, marker)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func countMessagesTokens(
|
||||
ctx context.Context,
|
||||
msgs []adk.Message,
|
||||
counter summarization.TokenCounterFunc,
|
||||
tools []*schema.ToolInfo,
|
||||
) (int, error) {
|
||||
if counter == nil {
|
||||
return 0, nil
|
||||
}
|
||||
n, err := counter(ctx, &summarization.TokenCounterInput{Messages: msgs, Tools: tools})
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return n, nil
|
||||
}
|
||||
|
||||
func truncateRoundMessagesToTokenBudget(
|
||||
ctx context.Context,
|
||||
round messageRound,
|
||||
tokenBudget int,
|
||||
counter summarization.TokenCounterFunc,
|
||||
toolMaxBytes int,
|
||||
spillRef string,
|
||||
) ([]adk.Message, error) {
|
||||
if tokenBudget <= 0 || len(round.messages) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
msgs := append([]adk.Message(nil), round.messages...)
|
||||
if n, err := countMessagesTokens(ctx, msgs, counter, nil); err != nil {
|
||||
return nil, err
|
||||
} else if n <= tokenBudget {
|
||||
return msgs, nil
|
||||
}
|
||||
if toolMaxBytes <= 0 {
|
||||
toolMaxBytes = 12000
|
||||
}
|
||||
for pass := 0; pass < 8 && toolMaxBytes >= 32; pass++ {
|
||||
out := make([]adk.Message, 0, len(msgs))
|
||||
for _, msg := range msgs {
|
||||
switch {
|
||||
case msg != nil && msg.Role == schema.Tool:
|
||||
out = append(out, truncateMessageToolContent(msg, toolMaxBytes, spillRef))
|
||||
case msg != nil && msg.Role == schema.Assistant:
|
||||
out = append(out, truncateMessageToolContent(msg, toolMaxBytes, spillRef))
|
||||
default:
|
||||
out = append(out, msg)
|
||||
}
|
||||
}
|
||||
n, err := countMessagesTokens(ctx, out, counter, nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if n <= tokenBudget {
|
||||
return out, nil
|
||||
}
|
||||
msgs = out
|
||||
toolMaxBytes /= 2
|
||||
}
|
||||
return msgs, nil
|
||||
}
|
||||
|
||||
type compactMessagesOpts struct {
|
||||
maxTokens int
|
||||
counter summarization.TokenCounterFunc
|
||||
toolMaxBytes int
|
||||
spillRef string
|
||||
aggressive bool
|
||||
logger *zap.Logger
|
||||
phase string
|
||||
}
|
||||
|
||||
func compactMessagesByDroppingRounds(
|
||||
ctx context.Context,
|
||||
messages []adk.Message,
|
||||
opts compactMessagesOpts,
|
||||
) ([]adk.Message, bool) {
|
||||
if opts.maxTokens <= 0 || len(messages) == 0 || opts.counter == nil {
|
||||
return messages, false
|
||||
}
|
||||
before, err := countMessagesTokens(ctx, messages, opts.counter, nil)
|
||||
if err != nil || before <= opts.maxTokens {
|
||||
return messages, false
|
||||
}
|
||||
|
||||
systems := make([]adk.Message, 0, 1)
|
||||
contextMsgs := make([]adk.Message, 0, len(messages))
|
||||
for _, msg := range messages {
|
||||
if msg != nil && msg.Role == schema.System && len(contextMsgs) == 0 {
|
||||
systems = append(systems, msg)
|
||||
continue
|
||||
}
|
||||
if msg != nil {
|
||||
contextMsgs = append(contextMsgs, msg)
|
||||
}
|
||||
}
|
||||
rounds := splitMessagesIntoRounds(contextMsgs)
|
||||
if len(rounds) == 0 {
|
||||
return messages, false
|
||||
}
|
||||
|
||||
startIdx := 0
|
||||
if opts.aggressive {
|
||||
startIdx = len(rounds) - 1
|
||||
if startIdx < 0 {
|
||||
startIdx = 0
|
||||
}
|
||||
}
|
||||
dropped := 0
|
||||
for len(rounds) > 1 || (opts.aggressive && len(rounds) == 1) {
|
||||
if !opts.aggressive && len(rounds) <= 1 {
|
||||
break
|
||||
}
|
||||
if opts.aggressive && len(rounds) == 1 {
|
||||
// Fall through to latest-round truncation below.
|
||||
break
|
||||
}
|
||||
rounds = rounds[1:]
|
||||
dropped++
|
||||
candidate := append([]adk.Message(nil), systems...)
|
||||
for _, round := range rounds {
|
||||
candidate = append(candidate, round.messages...)
|
||||
}
|
||||
after, countErr := countMessagesTokens(ctx, candidate, opts.counter, nil)
|
||||
if countErr != nil {
|
||||
break
|
||||
}
|
||||
if after <= opts.maxTokens {
|
||||
if opts.logger != nil {
|
||||
opts.logger.Warn("eino context compacted by dropping older rounds",
|
||||
zap.String("phase", opts.phase),
|
||||
zap.Int("tokens_before", before),
|
||||
zap.Int("tokens_after", after),
|
||||
zap.Int("max_tokens", opts.maxTokens),
|
||||
zap.Int("dropped_rounds", dropped),
|
||||
zap.Bool("aggressive", opts.aggressive),
|
||||
)
|
||||
}
|
||||
return candidate, true
|
||||
}
|
||||
if opts.aggressive {
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
if len(rounds) == 0 {
|
||||
return messages, false
|
||||
}
|
||||
latest := rounds[len(rounds)-1]
|
||||
truncated, truncErr := truncateRoundMessagesToTokenBudget(
|
||||
ctx, latest, opts.maxTokens, opts.counter, opts.toolMaxBytes, opts.spillRef,
|
||||
)
|
||||
if truncErr != nil || len(truncated) == 0 {
|
||||
if opts.logger != nil {
|
||||
opts.logger.Warn("eino context still above budget after round compaction; passing through without local error",
|
||||
zap.String("phase", opts.phase),
|
||||
zap.Int("tokens_before", before),
|
||||
zap.Int("max_tokens", opts.maxTokens),
|
||||
zap.Bool("aggressive", opts.aggressive),
|
||||
)
|
||||
}
|
||||
return messages, false
|
||||
}
|
||||
candidate := append([]adk.Message(nil), systems...)
|
||||
if dropped > 0 || startIdx > 0 {
|
||||
for _, round := range rounds[:len(rounds)-1] {
|
||||
candidate = append(candidate, round.messages...)
|
||||
}
|
||||
}
|
||||
candidate = append(candidate, truncated...)
|
||||
after, countErr := countMessagesTokens(ctx, candidate, opts.counter, nil)
|
||||
if countErr != nil {
|
||||
return messages, false
|
||||
}
|
||||
if opts.logger != nil {
|
||||
opts.logger.Warn("eino context compacted by truncating latest round tool output",
|
||||
zap.String("phase", opts.phase),
|
||||
zap.Int("tokens_before", before),
|
||||
zap.Int("tokens_after", after),
|
||||
zap.Int("max_tokens", opts.maxTokens),
|
||||
zap.Int("dropped_rounds", dropped),
|
||||
zap.Bool("aggressive", opts.aggressive),
|
||||
)
|
||||
}
|
||||
return candidate, true
|
||||
}
|
||||
|
||||
func aggressiveCompactMessagesForOverflow(
|
||||
ctx context.Context,
|
||||
messages []adk.Message,
|
||||
maxTotalTokens int,
|
||||
modelName string,
|
||||
toolMaxBytes int,
|
||||
phase string,
|
||||
logger *zap.Logger,
|
||||
) []adk.Message {
|
||||
if len(messages) == 0 || maxTotalTokens <= 0 {
|
||||
return messages
|
||||
}
|
||||
budget := maxTotalTokens * 70 / 100
|
||||
if budget < 4096 {
|
||||
budget = 4096
|
||||
}
|
||||
aggressiveToolMax := toolMaxBytes / aggressiveToolTruncDivisor
|
||||
if aggressiveToolMax < 2048 {
|
||||
aggressiveToolMax = 2048
|
||||
}
|
||||
out, _ := compactMessagesByDroppingRounds(ctx, messages, compactMessagesOpts{
|
||||
maxTokens: budget,
|
||||
counter: einoSummarizationTokenCounter(modelName),
|
||||
toolMaxBytes: aggressiveToolMax,
|
||||
aggressive: true,
|
||||
logger: logger,
|
||||
phase: phase,
|
||||
})
|
||||
return out
|
||||
}
|
||||
Reference in New Issue
Block a user