mirror of
https://github.com/alibaba/open-code-review.git
synced 2026-07-09 17:28:58 +00:00
feat: 记忆压缩三层分区 + 双阈值异步压
This commit is contained in:
kite
parent
8a93673268
commit
14adf6f55c
2 changed files with 289 additions and 58 deletions
|
|
@ -95,6 +95,33 @@ type AgentWarning struct {
|
|||
Type string `json:"type"`
|
||||
}
|
||||
|
||||
// compression thresholds as fractions of MaxTokens.
|
||||
const (
|
||||
tokenSoftThreshold = 0.60 // async background compression
|
||||
tokenWarningThreshold = 0.80 // immediate sync compression
|
||||
)
|
||||
|
||||
// round groups consecutive messages starting with an assistant message
|
||||
// followed by zero or more tool result messages.
|
||||
type round struct {
|
||||
assistantIdx int
|
||||
toolIdxs []int
|
||||
}
|
||||
|
||||
// partitionResult describes how messages should be split for compression.
|
||||
type partitionResult struct {
|
||||
frozenEnd int
|
||||
compressEnd int
|
||||
rounds []round
|
||||
activeCount int
|
||||
}
|
||||
|
||||
// compressionJob tracks an in-flight background compression operation.
|
||||
type compressionJob struct {
|
||||
done chan struct{}
|
||||
rebuilt []llm.Message
|
||||
}
|
||||
|
||||
// Agent orchestrates the AI-powered code review.
|
||||
type Agent struct {
|
||||
args Args
|
||||
|
|
@ -108,6 +135,8 @@ type Agent struct {
|
|||
totalOutputTokens int64 // accumulated completion tokens, accessed atomically
|
||||
warningsMu sync.Mutex
|
||||
warnings []AgentWarning
|
||||
compressionMu sync.Mutex
|
||||
pendingJob *compressionJob
|
||||
}
|
||||
|
||||
// CommentWorkerPool manages a fixed-size pool of workers dedicated to
|
||||
|
|
@ -812,27 +841,51 @@ func (a *Agent) collectPendingComments() []model.LlmComment {
|
|||
}
|
||||
|
||||
// addNextMessage adds assistant + tool response messages to the conversation history.
|
||||
// Implements dual-threshold compression:
|
||||
// - 60% of MaxTokens: trigger async background compression (non-blocking)
|
||||
// - 80% of MaxTokens: perform synchronous compression immediately
|
||||
func (a *Agent) addNextMessage(assistantContent string, toolCalls []llm.ToolCall, results []tool.ToolCallResult, messages *[]llm.Message, filePath string) bool {
|
||||
// Check if context compression is needed (80% of MaxTokens)
|
||||
maxAllowed := a.args.Template.MaxTokens
|
||||
softLimit := int(float64(maxAllowed) * tokenSoftThreshold)
|
||||
warnLimit := int(float64(maxAllowed) * tokenWarningThreshold)
|
||||
|
||||
// Try to apply any completed async compression from a previous iteration.
|
||||
a.tryApplyPendingCompression(messages)
|
||||
|
||||
tokenCount := countMessagesTokens(*messages)
|
||||
limit := a.args.Template.MaxTokens * 4 / 5
|
||||
if tokenCount > limit {
|
||||
*messages = a.compressAndRecord(*messages, filePath)
|
||||
|
||||
// Hard threshold: synchronous compression.
|
||||
if tokenCount > warnLimit {
|
||||
a.cancelPendingCompression()
|
||||
*messages, _ = a.runCompression(*messages, filePath)
|
||||
tokenCount = countMessagesTokens(*messages)
|
||||
}
|
||||
|
||||
// Add assistant message with tool_calls when present
|
||||
// Soft threshold: async compression.
|
||||
if tokenCount > softLimit && a.pendingJob == nil {
|
||||
a.triggerAsyncCompression(*messages, filePath)
|
||||
}
|
||||
|
||||
// Add assistant message with tool_calls when present.
|
||||
if len(toolCalls) > 0 {
|
||||
*messages = append(*messages, llm.NewToolCallMessage(assistantContent, toolCalls))
|
||||
} else if assistantContent != "" {
|
||||
*messages = append(*messages, llm.NewTextMessage("assistant", assistantContent))
|
||||
}
|
||||
|
||||
// Add tool response messages using Claude's tool_result format
|
||||
// Add tool response messages using Claude's tool_result format.
|
||||
for _, r := range results {
|
||||
*messages = append(*messages, llm.NewToolResultMessage(r.ToolCallID, r.Result))
|
||||
}
|
||||
|
||||
return countMessagesTokens(*messages) < limit
|
||||
// Final check: compress synchronously if still over warning limit.
|
||||
finalCount := countMessagesTokens(*messages)
|
||||
if finalCount > warnLimit {
|
||||
a.cancelPendingCompression()
|
||||
*messages, _ = a.runCompression(*messages, filePath)
|
||||
}
|
||||
|
||||
return countMessagesTokens(*messages) < warnLimit
|
||||
}
|
||||
|
||||
func countMessagesTokens(msgs []llm.Message) int {
|
||||
|
|
@ -843,6 +896,233 @@ func countMessagesTokens(msgs []llm.Message) int {
|
|||
return total
|
||||
}
|
||||
|
||||
// groupIntoRounds parses messages[start:] into logical (assistant + tool_results) pairs.
|
||||
func groupIntoRounds(messages []llm.Message, start int) []round {
|
||||
var rounds []round
|
||||
i := start
|
||||
for i < len(messages) {
|
||||
if messages[i].Role == "assistant" {
|
||||
r := round{assistantIdx: i}
|
||||
i++
|
||||
for i < len(messages) && messages[i].Role == "tool" {
|
||||
r.toolIdxs = append(r.toolIdxs, i)
|
||||
i++
|
||||
}
|
||||
rounds = append(rounds, r)
|
||||
} else {
|
||||
i++
|
||||
}
|
||||
}
|
||||
return rounds
|
||||
}
|
||||
|
||||
// computeActiveZoneSize returns how many trailing rounds fit within the remaining
|
||||
// token budget after accounting for frozen zone and the compressed summary.
|
||||
func computeActiveZoneSize(rounds []round, messages []llm.Message, maxTokens int, reservedTokens int) int {
|
||||
budget := int(float64(maxTokens)*tokenWarningThreshold) - reservedTokens
|
||||
if budget <= 0 {
|
||||
return 0
|
||||
}
|
||||
|
||||
count := 0
|
||||
tokensUsed := 0
|
||||
for i := len(rounds) - 1; i >= 0; i-- {
|
||||
roundTokens := llm.CountTokens(messages[rounds[i].assistantIdx].ExtractText())
|
||||
for _, ti := range rounds[i].toolIdxs {
|
||||
roundTokens += llm.CountTokens(messages[ti].ExtractText())
|
||||
}
|
||||
if tokensUsed+roundTokens > budget {
|
||||
break
|
||||
}
|
||||
tokensUsed += roundTokens
|
||||
count++
|
||||
}
|
||||
return count
|
||||
}
|
||||
|
||||
// partitionMessages divides messages into frozen, compress, and active zones.
|
||||
// Frozen zone is always messages[0:2]. Active zone preserves the K most recent
|
||||
// complete rounds based on available token budget.
|
||||
func partitionMessages(messages []llm.Message, maxTokens int, prevSummaryTokenEstimate int) partitionResult {
|
||||
result := partitionResult{frozenEnd: 2}
|
||||
if len(messages) <= 2 {
|
||||
result.compressEnd = len(messages)
|
||||
return result
|
||||
}
|
||||
|
||||
result.rounds = groupIntoRounds(messages, 2)
|
||||
if len(result.rounds) == 0 {
|
||||
result.compressEnd = len(messages)
|
||||
return result
|
||||
}
|
||||
|
||||
result.activeCount = computeActiveZoneSize(result.rounds, messages, maxTokens, prevSummaryTokenEstimate)
|
||||
if result.activeCount >= len(result.rounds) {
|
||||
// Everything fits — no compression needed.
|
||||
result.compressEnd = len(messages)
|
||||
result.activeCount = 0
|
||||
return result
|
||||
}
|
||||
|
||||
// compressEnd = index after the last round NOT in active zone.
|
||||
activeStartIdx := len(result.rounds) - result.activeCount
|
||||
lastCompressRound := result.rounds[activeStartIdx-1]
|
||||
if len(lastCompressRound.toolIdxs) > 0 {
|
||||
result.compressEnd = lastCompressRound.toolIdxs[len(lastCompressRound.toolIdxs)-1] + 1
|
||||
} else {
|
||||
result.compressEnd = lastCompressRound.assistantIdx + 1
|
||||
}
|
||||
|
||||
return result
|
||||
}
|
||||
|
||||
// stripMarkdownFences removes ```json and ``` wrappers that some models
|
||||
// add around structured outputs.
|
||||
func stripMarkdownFences(s string) string {
|
||||
s = strings.TrimSpace(s)
|
||||
if strings.HasPrefix(s, "```") {
|
||||
if nl := strings.IndexByte(s, '\n'); nl >= 0 {
|
||||
s = s[nl+1:]
|
||||
} else {
|
||||
s = strings.TrimPrefix(s, "```json")
|
||||
s = strings.TrimPrefix(s, "```")
|
||||
}
|
||||
}
|
||||
s = strings.TrimSpace(s)
|
||||
if strings.HasSuffix(s, "```") {
|
||||
s = strings.TrimSuffix(s, "```")
|
||||
s = strings.TrimSpace(s)
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
// runCompression performs three-zone memory compression on the given messages.
|
||||
// It summarizes the compress zone while preserving the active zone intact.
|
||||
// Returns the rebuilt messages slice: [frozen] + [compressed_summary] + [active].
|
||||
func (a *Agent) runCompression(msgs []llm.Message, filePath string) ([]llm.Message, error) {
|
||||
if len(a.args.Template.MemoryCompressionTask.Messages) == 0 || len(msgs) <= 2 {
|
||||
return msgs[:min(len(msgs), 2)], nil
|
||||
}
|
||||
|
||||
part := partitionMessages(msgs, a.args.Template.MaxTokens, 0)
|
||||
if part.compressEnd <= part.frozenEnd {
|
||||
return msgs, nil
|
||||
}
|
||||
|
||||
contextXML := buildMessageXML(msgs[part.frozenEnd:part.compressEnd])
|
||||
|
||||
compressionMsgs := make([]llm.Message, 0, len(a.args.Template.MemoryCompressionTask.Messages))
|
||||
for _, m := range a.args.Template.MemoryCompressionTask.Messages {
|
||||
content := strings.ReplaceAll(m.Content, "{{context}}", contextXML)
|
||||
compressionMsgs = append(compressionMsgs, llm.NewTextMessage(m.Role, content))
|
||||
}
|
||||
|
||||
startTime := time.Now()
|
||||
resp, err := a.args.LLMClient.Completions(llm.ChatRequest{
|
||||
Model: a.args.Model,
|
||||
Messages: compressionMsgs,
|
||||
MaxTokens: a.args.Template.MaxTokens,
|
||||
})
|
||||
duration := time.Since(startTime)
|
||||
|
||||
fs := a.session.GetOrCreateFileSession(filePath)
|
||||
rec := fs.AppendTaskRecord(session.MemoryCompressionTask, compressionMsgs)
|
||||
if err != nil {
|
||||
rec.SetError(err, duration)
|
||||
fmt.Fprintf(stdout.Writer(), "[ocr] Memory compression failed: %v\n", err)
|
||||
return msgs[:part.frozenEnd], fmt.Errorf("memory compression: %w", err)
|
||||
}
|
||||
rec.SetResponse(resp, duration)
|
||||
if resp.Usage != nil {
|
||||
atomic.AddInt64(&a.totalTokensUsed, int64(resp.Usage.TotalTokens))
|
||||
atomic.AddInt64(&a.totalInputTokens, int64(resp.Usage.PromptTokens+resp.Usage.CacheReadTokens))
|
||||
atomic.AddInt64(&a.totalOutputTokens, int64(resp.Usage.CompletionTokens+resp.Usage.CacheWriteTokens))
|
||||
}
|
||||
|
||||
rawSummary := stripMarkdownFences(resp.Content())
|
||||
if rawSummary == "" {
|
||||
return msgs[:part.frozenEnd], nil
|
||||
}
|
||||
|
||||
rebuilt := make([]llm.Message, 2)
|
||||
copy(rebuilt, msgs[:2])
|
||||
|
||||
userMsg := rebuilt[1]
|
||||
currentText := userMsg.ExtractText()
|
||||
rebuilt[1] = llm.NewTextMessage(userMsg.Role, currentText+"\n\n<previous_review_summary>\n"+rawSummary+"\n</previous_review_summary>")
|
||||
|
||||
for i := part.compressEnd; i < len(msgs); i++ {
|
||||
rebuilt = append(rebuilt, msgs[i])
|
||||
}
|
||||
|
||||
return rebuilt, nil
|
||||
}
|
||||
|
||||
// triggerAsyncCompression kicks off a background compression job.
|
||||
func (a *Agent) triggerAsyncCompression(messages []llm.Message, filePath string) {
|
||||
msgSnapshot := copyMessages(messages)
|
||||
|
||||
job := &compressionJob{done: make(chan struct{})}
|
||||
a.compressionMu.Lock()
|
||||
a.pendingJob = job
|
||||
a.compressionMu.Unlock()
|
||||
|
||||
go func() {
|
||||
rebuilt, _ := a.runCompression(msgSnapshot, filePath)
|
||||
|
||||
a.compressionMu.Lock()
|
||||
defer a.compressionMu.Unlock()
|
||||
|
||||
if a.pendingJob != job {
|
||||
return // cancelled or superseded
|
||||
}
|
||||
job.rebuilt = rebuilt
|
||||
close(job.done)
|
||||
}()
|
||||
}
|
||||
|
||||
// tryApplyPendingCompression checks if a background compression completed
|
||||
// and swaps the rebuilt messages into place. Returns true if applied.
|
||||
func (a *Agent) tryApplyPendingCompression(messages *[]llm.Message) bool {
|
||||
a.compressionMu.Lock()
|
||||
job := a.pendingJob
|
||||
a.compressionMu.Unlock()
|
||||
|
||||
if job == nil {
|
||||
return false
|
||||
}
|
||||
|
||||
select {
|
||||
case <-job.done:
|
||||
a.compressionMu.Lock()
|
||||
if a.pendingJob == job && job.rebuilt != nil {
|
||||
*messages = job.rebuilt
|
||||
a.pendingJob = nil
|
||||
}
|
||||
a.compressionMu.Unlock()
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
// cancelPendingCompression aborts any in-flight background compression.
|
||||
func (a *Agent) cancelPendingCompression() {
|
||||
a.compressionMu.Lock()
|
||||
defer a.compressionMu.Unlock()
|
||||
|
||||
if a.pendingJob != nil {
|
||||
a.pendingJob = nil
|
||||
}
|
||||
}
|
||||
|
||||
// copyMessages creates a shallow copy of a message slice.
|
||||
func copyMessages(msgs []llm.Message) []llm.Message {
|
||||
out := make([]llm.Message, len(msgs))
|
||||
copy(out, msgs)
|
||||
return out
|
||||
}
|
||||
|
||||
// BuildToolDefs converts toolsconfig.ToolConfigEntry slice into []llm.ToolDef,
|
||||
// filtering by phase (planOnly=true for plan_task, false for main_task).
|
||||
func BuildToolDefs(entries []toolsconfig.ToolConfigEntry, planOnly bool) []llm.ToolDef {
|
||||
|
|
@ -865,55 +1145,6 @@ func BuildToolDefs(entries []toolsconfig.ToolConfigEntry, planOnly bool) []llm.T
|
|||
return defs
|
||||
}
|
||||
|
||||
// compressAndRecord runs memory compression and records it in session history.
|
||||
func (a *Agent) compressAndRecord(msgs []llm.Message, filePath string) []llm.Message {
|
||||
if len(a.args.Template.MemoryCompressionTask.Messages) == 0 || len(msgs) <= 2 {
|
||||
return msgs[:min(len(msgs), 2)]
|
||||
}
|
||||
|
||||
contextXML := buildMessageXML(msgs[2:])
|
||||
compressionMsgs := make([]llm.Message, 0, len(a.args.Template.MemoryCompressionTask.Messages))
|
||||
for _, m := range a.args.Template.MemoryCompressionTask.Messages {
|
||||
content := strings.ReplaceAll(m.Content, "{{context}}", contextXML)
|
||||
compressionMsgs = append(compressionMsgs, llm.NewTextMessage(m.Role, content))
|
||||
}
|
||||
|
||||
startTime := time.Now()
|
||||
resp, err := a.args.LLMClient.Completions(llm.ChatRequest{
|
||||
Model: a.args.Model,
|
||||
Messages: compressionMsgs,
|
||||
MaxTokens: a.args.Template.MaxTokens,
|
||||
})
|
||||
duration := time.Since(startTime)
|
||||
|
||||
fs := a.session.GetOrCreateFileSession(filePath)
|
||||
rec := fs.AppendTaskRecord(session.MemoryCompressionTask, compressionMsgs)
|
||||
if err != nil {
|
||||
rec.SetError(err, duration)
|
||||
fmt.Fprintf(stdout.Writer(), "[ocr] Memory compression failed: %v\n", err)
|
||||
return msgs[:2]
|
||||
}
|
||||
rec.SetResponse(resp, duration)
|
||||
if resp.Usage != nil {
|
||||
atomic.AddInt64(&a.totalTokensUsed, int64(resp.Usage.TotalTokens))
|
||||
atomic.AddInt64(&a.totalInputTokens, int64(resp.Usage.PromptTokens+resp.Usage.CacheReadTokens))
|
||||
atomic.AddInt64(&a.totalOutputTokens, int64(resp.Usage.CompletionTokens+resp.Usage.CacheWriteTokens))
|
||||
}
|
||||
|
||||
summary := resp.Content()
|
||||
if summary == "" {
|
||||
return msgs[:2]
|
||||
}
|
||||
|
||||
compressed := make([]llm.Message, 2)
|
||||
copy(compressed, msgs[:2])
|
||||
// Append summary to the original user prompt
|
||||
userMsg := compressed[1]
|
||||
currentText := userMsg.ExtractText()
|
||||
compressed[1] = llm.NewTextMessage(userMsg.Role, currentText+"\n\n"+summary)
|
||||
return compressed
|
||||
}
|
||||
|
||||
func buildMessageXML(msgs []llm.Message) string {
|
||||
var sb strings.Builder
|
||||
for i, m := range msgs {
|
||||
|
|
|
|||
|
|
@ -29,7 +29,7 @@
|
|||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
"content": "<目标>\n你是一个专业的对话总结助手。你会收到一段用户和助手之间的对话。压缩对话历史记录，同时保持执行的连续性，使代码评审助手能够从当前状态继续执行任务，而无需重新启动。\n</目标>\n\n对话历史记录了代码评审助手在帮助用户审查代码时调用工具的原因以及调用工具时返回的上下文。\n\n\n<正面标准>\n应该包括以下内容：\n- 由于xxx原因，我使用了xxx工具，查看了xxx上下文。\n- 查看xxx上下文后，由于xxx原因，我使用了xxx工具。\n- ...\n</正面标准>\n\n<负面标准>\n不要包括：\n- 上下文中的具体细节\n- 以后不会相关的临时上下文\n</负面标准>\n\n<返回示例>\n<previous_summary>\n- xxx。\n- yyy。\n- ......\n</previous_summary>\n<返回示例>\n\n请你直接返回答案，不要返回其他无关内容。"
|
||||
"content": "## Goal\nYou are a professional code review conversation summarization assistant. You will receive a conversation history between a code review assistant and an LLM model (including tool calls and their results). Compress this conversation into a structured summary so that the code review assistant can continue from the current state without restarting.\n\n## Output Format Requirements\nOrganize the summary using the following five dimensions, separated by explicit headings:\n\n### Identified Code Issues\nList all confirmed issues sorted by severity (HIGH / MEDIUM / LOW). Each entry should include: file path, issue type, severity, brief description. Example:\n- [HIGH] `UserService.go:45` — map concurrent read-write access without lock, suggest adding sync.RWMutex\n- [MEDIUM] `config_loader.go:12` — error handling is incomplete, may swallow critical information\n\n### Tool Call Conclusions\nSummarize key findings and conclusions from each tool invocation. Example:\n- get_function_info(UserService): confirmed concurrent write-to-map logic within this function\n- search_file(\"database\"): no other related configuration issues found\n\n### Completed Tasks\nList items that have been completed and require no further follow-up.\n\n### Pending Tasks\nList items that have been started but not yet completed and still need attention.\n\n### Current Focus\nSummarize in one sentence the core matter currently being investigated or handled.\n\n## Rules\n1. Do not include specific code details; only reference file paths and issue types\n2. Avoid repetitive or redundant information\n3. Omit any dimension that has no relevant content\n4. Completed/pending task list items should be described as complete sentences\n5. current_focus should be concise, no more than one sentence"
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
|
|
@ -43,4 +43,4 @@
|
|||
"MAX_SUBTASK_EXECUTION_TIME_MINUTES": 5,
|
||||
"PLAN_MODE_LINE_THRESHOLD": 50,
|
||||
"MAX_TOKENS": 58888
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue