98 lines
3.1 KiB
Go
98 lines
3.1 KiB
Go
package compaction
|
|
|
|
import "github.com/smallnest/pigo/internal/agentcore"
|
|
|
|
// CutPointResult describes the cut selected for compaction, mirroring pi's
|
|
// CutPointResult (adapted to pigo's flat message list).
|
|
type CutPointResult struct {
|
|
// FirstKeptIndex is the index of the first message retained after
|
|
// compaction; everything before it is summarized.
|
|
FirstKeptIndex int
|
|
// TurnStartIndex is the index of the user message that starts the turn the
|
|
// cut falls inside, or -1 when the cut lands on a clean turn boundary.
|
|
TurnStartIndex int
|
|
// IsSplitTurn reports whether the cut splits an in-progress assistant turn.
|
|
IsSplitTurn bool
|
|
}
|
|
|
|
// isValidCutPoint reports whether a message may serve as a cut point. A cut may
|
|
// land on a user or assistant message but never on a toolResult, because a
|
|
// toolResult must stay attached to the toolCall that produced it (pi semantics).
|
|
func isValidCutPoint(msg agentcore.Message) bool {
|
|
switch msg.Role() {
|
|
case agentcore.RoleUser, agentcore.RoleAssistant:
|
|
return true
|
|
default: // toolResult and any custom kinds are not cuttable.
|
|
return false
|
|
}
|
|
}
|
|
|
|
// findValidCutPoints returns the indices of all messages that may serve as cut
|
|
// points, in ascending order.
|
|
func findValidCutPoints(msgs []agentcore.Message) []int {
|
|
var pts []int
|
|
for i, m := range msgs {
|
|
if isValidCutPoint(m) {
|
|
pts = append(pts, i)
|
|
}
|
|
}
|
|
return pts
|
|
}
|
|
|
|
// findTurnStartIndex scans backwards from idx to find the user message that
|
|
// starts the turn containing idx, returning -1 if none is found.
|
|
func findTurnStartIndex(msgs []agentcore.Message, idx int) int {
|
|
for i := idx; i >= 0; i-- {
|
|
if msgs[i].Role() == agentcore.RoleUser {
|
|
return i
|
|
}
|
|
}
|
|
return -1
|
|
}
|
|
|
|
// FindCutPoint finds the compaction cut that keeps approximately
|
|
// keepRecentTokens worth of the most recent messages. It accumulates token
|
|
// estimates from the newest message backwards; once the retained budget is
|
|
// reached it snaps to the nearest valid cut point at or after that message,
|
|
// never splitting a toolCall from its toolResult. This mirrors pi's findCutPoint.
|
|
//
|
|
// When no valid cut point exists, it keeps everything (FirstKeptIndex 0).
|
|
func FindCutPoint(msgs []agentcore.Message, keepRecentTokens int) CutPointResult {
|
|
cutPoints := findValidCutPoints(msgs)
|
|
if len(cutPoints) == 0 {
|
|
return CutPointResult{FirstKeptIndex: 0, TurnStartIndex: -1, IsSplitTurn: false}
|
|
}
|
|
|
|
// Default to the earliest valid cut point (keep as much as possible) when
|
|
// the retained budget is never reached.
|
|
cutIndex := cutPoints[0]
|
|
|
|
accumulated := 0
|
|
for i := len(msgs) - 1; i >= 0; i-- {
|
|
accumulated += EstimateTokens(msgs[i])
|
|
if accumulated >= keepRecentTokens {
|
|
// Snap to the nearest valid cut point at or after i, so the kept
|
|
// window starts on a cuttable boundary.
|
|
for _, c := range cutPoints {
|
|
if c >= i {
|
|
cutIndex = c
|
|
break
|
|
}
|
|
}
|
|
break
|
|
}
|
|
}
|
|
|
|
cutOnUser := msgs[cutIndex].Role() == agentcore.RoleUser
|
|
turnStart := -1
|
|
if !cutOnUser {
|
|
turnStart = findTurnStartIndex(msgs, cutIndex)
|
|
}
|
|
|
|
return CutPointResult{
|
|
FirstKeptIndex: cutIndex,
|
|
TurnStartIndex: turnStart,
|
|
IsSplitTurn: !cutOnUser && turnStart != -1,
|
|
}
|
|
}
|