Add files via upload

internal/multiagent/model_input_soft_budget_middleware.go

@@ -0,0 +1,66 @@
+package multiagent
+
+import (
+	"context"
+
+	"github.com/cloudwego/eino/adk"
+	"github.com/cloudwego/eino/adk/middlewares/summarization"
+	"go.uber.org/zap"
+)
+
+// modelInputSoftBudgetMiddleware is the final guard before a normal model call.
+// It drops oldest complete rounds and truncates oversized tool output in the latest
+// round, but never fails locally — API context limits are handled by overflow retry.
+type modelInputSoftBudgetMiddleware struct {
+	adk.BaseChatModelAgentMiddleware
+	maxTokens    int
+	toolMaxBytes int
+	counter      summarization.TokenCounterFunc
+	logger       *zap.Logger
+	phase        string
+}
+
+func newModelInputSoftBudgetMiddleware(
+	maxTotalTokens int,
+	toolMaxBytes int,
+	modelName string,
+	logger *zap.Logger,
+	phase string,
+) adk.ChatModelAgentMiddleware {
+	if maxTotalTokens <= 0 {
+		maxTotalTokens = 120000
+	}
+	if toolMaxBytes <= 0 {
+		toolMaxBytes = 12000
+	}
+	return &modelInputSoftBudgetMiddleware{
+		maxTokens:    maxTotalTokens,
+		toolMaxBytes: toolMaxBytes,
+		counter:      einoSummarizationTokenCounter(modelName),
+		logger:       logger,
+		phase:        phase,
+	}
+}
+
+func (m *modelInputSoftBudgetMiddleware) BeforeModelRewriteState(
+	ctx context.Context,
+	state *adk.ChatModelAgentState,
+	mc *adk.ModelContext,
+) (context.Context, *adk.ChatModelAgentState, error) {
+	if m == nil || state == nil || len(state.Messages) == 0 {
+		return ctx, state, nil
+	}
+	compacted, changed := compactMessagesByDroppingRounds(ctx, state.Messages, compactMessagesOpts{
+		maxTokens:    m.maxTokens,
+		counter:      m.counter,
+		toolMaxBytes: m.toolMaxBytes,
+		phase:        m.phase,
+		logger:       m.logger,
+	})
+	if !changed {
+		return ctx, state, nil
+	}
+	out := *state
+	out.Messages = compacted
+	return ctx, &out, nil
+}