109 lines
2.9 KiB
Go
109 lines
2.9 KiB
Go
package ctxmgr
|
|
|
|
import (
|
|
"fmt"
|
|
"strings"
|
|
)
|
|
|
|
// WindowConfig 滑动窗口配置
|
|
type WindowConfig struct {
|
|
MaxTurns int // 窗口内最大消息条数（一问一答计 2 条），默认 16
|
|
TokenBudget int // 窗口内字符/Token 预算，默认 3000
|
|
ReserveTail int // 为本轮用户输入与固定 Prompt 预留的字符预算，默认 800
|
|
}
|
|
|
|
// DefaultWindowConfig 返回默认滑动窗口配置
|
|
func DefaultWindowConfig() WindowConfig {
|
|
return WindowConfig{
|
|
MaxTurns: 16,
|
|
TokenBudget: 3000,
|
|
ReserveTail: 800,
|
|
}
|
|
}
|
|
|
|
// EstimateTokens 估算文本 Token 数 (中文及中英混排场景下采用保守 rune 计数：1 rune ≈ 1 token)
|
|
func EstimateTokens(s string) int {
|
|
return len([]rune(s))
|
|
}
|
|
|
|
// AssembleWindow 从完整历史消息列表中按双预算（条数上限 + Token预算）逆序裁剪出近期对话窗口
|
|
// 返回：
|
|
// - window: 窗口内的近期消息（按原始时间正序）
|
|
// - early: 窗口之外被淘汰的早期消息（供滚动摘要器提取与压缩）
|
|
// - truncated: 是否发生了截断
|
|
func AssembleWindow(msgs []ChatMessageRecord, cfg WindowConfig) (window []ChatMessageRecord, early []ChatMessageRecord, truncated bool) {
|
|
if len(msgs) == 0 {
|
|
return []ChatMessageRecord{}, []ChatMessageRecord{}, false
|
|
}
|
|
|
|
if cfg.MaxTurns <= 0 {
|
|
cfg.MaxTurns = 16
|
|
}
|
|
if cfg.TokenBudget <= 0 {
|
|
cfg.TokenBudget = 3000
|
|
}
|
|
|
|
netBudget := cfg.TokenBudget - cfg.ReserveTail
|
|
if netBudget < 400 {
|
|
netBudget = cfg.TokenBudget // 若预留过大，以全额预算为准
|
|
}
|
|
|
|
totalTokens := 0
|
|
cutIndex := 0
|
|
|
|
// 逆序向前收纳
|
|
for i := len(msgs) - 1; i >= 0; i-- {
|
|
m := msgs[i]
|
|
tokens := EstimateTokens(m.Content)
|
|
if tokens == 0 {
|
|
tokens = 1
|
|
}
|
|
|
|
// 检查条数预算与 Token 预算
|
|
countInWindow := len(msgs) - i
|
|
if countInWindow > cfg.MaxTurns || (totalTokens+tokens > netBudget && countInWindow > 2) {
|
|
cutIndex = i + 1
|
|
break
|
|
}
|
|
|
|
totalTokens += tokens
|
|
}
|
|
|
|
if cutIndex > 0 {
|
|
early = make([]ChatMessageRecord, cutIndex)
|
|
copy(early, msgs[:cutIndex])
|
|
window = make([]ChatMessageRecord, len(msgs)-cutIndex)
|
|
copy(window, msgs[cutIndex:])
|
|
return window, early, true
|
|
}
|
|
|
|
window = make([]ChatMessageRecord, len(msgs))
|
|
copy(window, msgs)
|
|
return window, []ChatMessageRecord{}, false
|
|
}
|
|
|
|
// FormatNumberedDialogue 将消息序列格式化为带逐轮序号的近期对话文本
|
|
// 格式如：
|
|
// [1] 用户: 帮我看看项目进展
|
|
// [2] 助手: 当前处于 V0.1 阶段...
|
|
func FormatNumberedDialogue(msgs []ChatMessageRecord, startIndex int) string {
|
|
if len(msgs) == 0 {
|
|
return ""
|
|
}
|
|
if startIndex <= 0 {
|
|
startIndex = 1
|
|
}
|
|
|
|
var sb strings.Builder
|
|
for i, m := range msgs {
|
|
roleName := "用户"
|
|
if m.Role == "assistant" {
|
|
roleName = "助手"
|
|
} else if m.Role == "system" {
|
|
roleName = "系统"
|
|
}
|
|
sb.WriteString(fmt.Sprintf("[%d] %s: %s\n", startIndex+i, roleName, strings.TrimSpace(m.Content)))
|
|
}
|
|
return sb.String()
|
|
}
|