96 lines
3.0 KiB
Go
96 lines
3.0 KiB
Go
package agent
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"strings"
|
|
|
|
"ai-agent/internal/llm"
|
|
)
|
|
|
|
const compactThreshold = 0.75
|
|
const keepMessages = 4
|
|
|
|
func (a *Agent) shouldCompact(promptTokens int) bool {
|
|
if a.numCtx <= 0 || promptTokens <= 0 {
|
|
return false
|
|
}
|
|
return float64(promptTokens) > float64(a.numCtx)*compactThreshold
|
|
}
|
|
|
|
func (a *Agent) compact(ctx context.Context, out Output) bool {
|
|
a.mu.RLock()
|
|
msgCount := len(a.messages)
|
|
a.mu.RUnlock()
|
|
if msgCount <= keepMessages+1 {
|
|
return false
|
|
}
|
|
a.mu.RLock()
|
|
splitAt := msgCount - keepMessages
|
|
older := make([]llm.Message, splitAt)
|
|
copy(older, a.messages[:splitAt])
|
|
recent := make([]llm.Message, keepMessages)
|
|
copy(recent, a.messages[splitAt:])
|
|
a.mu.RUnlock()
|
|
summary := summarizeMessages(older)
|
|
var summaryBuf strings.Builder
|
|
err := a.llmClient.ChatStream(ctx, llm.ChatOptions{
|
|
Messages: []llm.Message{
|
|
{Role: "user", Content: summary},
|
|
},
|
|
System: "You are a conversation summarizer. Produce a concise summary of the conversation so far, capturing all key facts, decisions, tool results, and user requests. Keep it under 500 words. Output only the summary, no preamble.",
|
|
}, func(chunk llm.StreamChunk) error {
|
|
if chunk.Text != "" {
|
|
summaryBuf.WriteString(chunk.Text)
|
|
}
|
|
return nil
|
|
})
|
|
if err != nil {
|
|
out.Error(fmt.Sprintf("compaction failed: %v", err))
|
|
return false
|
|
}
|
|
summaryText := summaryBuf.String()
|
|
if summaryText == "" {
|
|
return false
|
|
}
|
|
if a.iceEngine != nil {
|
|
if err := a.iceEngine.IndexSummary(ctx, summaryText); err != nil {
|
|
out.Error(fmt.Sprintf("ICE summary indexing failed: %v", err))
|
|
}
|
|
}
|
|
compacted := make([]llm.Message, 0, 1+len(recent))
|
|
compacted = append(compacted, llm.Message{
|
|
Role: "user",
|
|
Content: fmt.Sprintf("[Conversation summary: %s]", summaryText),
|
|
})
|
|
compacted = append(compacted, recent...)
|
|
a.ReplaceMessages(compacted)
|
|
out.SystemMessage(fmt.Sprintf("Context compacted: %d messages summarized, %d kept", len(older), len(recent)))
|
|
return true
|
|
}
|
|
|
|
func summarizeMessages(msgs []llm.Message) string {
|
|
var b strings.Builder
|
|
b.WriteString("Summarize this conversation:\n\n")
|
|
for _, msg := range msgs {
|
|
switch msg.Role {
|
|
case "user":
|
|
fmt.Fprintf(&b, "User: %s\n", msg.Content)
|
|
case "assistant":
|
|
if msg.Content != "" {
|
|
fmt.Fprintf(&b, "Assistant: %s\n", msg.Content)
|
|
}
|
|
for _, tc := range msg.ToolCalls {
|
|
fmt.Fprintf(&b, "Assistant called tool %s(%s)\n", tc.Name, FormatToolArgs(tc.Arguments))
|
|
}
|
|
case "tool":
|
|
content := msg.Content
|
|
if len(content) > 300 {
|
|
content = content[:297] + "..."
|
|
}
|
|
fmt.Fprintf(&b, "Tool %s result: %s\n", msg.ToolName, content)
|
|
}
|
|
}
|
|
return b.String()
|
|
}
|