From 07e624c8dacc6bf630298e4a55fb9d9d8f8e62b5 Mon Sep 17 00:00:00 2001 From: Danieldd28 Date: Tue, 10 Feb 2026 01:25:46 +0700 Subject: [PATCH] feat: implement dynamic context compression for efficient memory usage - Added Summary field to Session struct - Implemented background summarization when history > 20 messages - Included conversation summary in system prompt for long-term context - Added thread-safety for concurrent summarization per session --- pkg/agent/context.go | 6 +++- pkg/agent/loop.go | 76 +++++++++++++++++++++++++++++++++++++++++- pkg/session/manager.go | 40 ++++++++++++++++++++++ 3 files changed, 120 insertions(+), 2 deletions(-) diff --git a/pkg/agent/context.go b/pkg/agent/context.go index 5a1a734..9ed5733 100644 --- a/pkg/agent/context.go +++ b/pkg/agent/context.go @@ -84,7 +84,7 @@ func (cb *ContextBuilder) LoadBootstrapFiles() string { return result } -func (cb *ContextBuilder) BuildMessages(history []providers.Message, currentMessage string, media []string) []providers.Message { +func (cb *ContextBuilder) BuildMessages(history []providers.Message, summary string, currentMessage string, media []string) []providers.Message { messages := []providers.Message{} systemPrompt := cb.BuildSystemPrompt() @@ -103,6 +103,10 @@ func (cb *ContextBuilder) BuildMessages(history []providers.Message, currentMess systemPrompt += "\n\n" + skillsContent } + if summary != "" { + systemPrompt += "\n\n## Summary of Previous Conversation\n\n" + summary + } + messages = append(messages, providers.Message{ Role: "system", Content: systemPrompt, diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index e23f4d1..05c8dc1 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -12,6 +12,8 @@ import ( "fmt" "os" "path/filepath" + "sync" + "time" "github.com/sipeed/picoclaw/pkg/bus" "github.com/sipeed/picoclaw/pkg/config" @@ -30,6 +32,7 @@ type AgentLoop struct { contextBuilder *ContextBuilder tools *tools.ToolRegistry running bool + summarizing sync.Map } func NewAgentLoop(cfg *config.Config, bus *bus.MessageBus, provider providers.LLMProvider) *AgentLoop { @@ -58,6 +61,7 @@ func NewAgentLoop(cfg *config.Config, bus *bus.MessageBus, provider providers.LL contextBuilder: NewContextBuilder(workspace), tools: toolsRegistry, running: false, + summarizing: sync.Map{}, } } @@ -109,8 +113,12 @@ func (al *AgentLoop) ProcessDirect(ctx context.Context, content, sessionKey stri } func (al *AgentLoop) processMessage(ctx context.Context, msg bus.InboundMessage) (string, error) { + history := al.sessions.GetHistory(msg.SessionKey) + summary := al.sessions.GetSummary(msg.SessionKey) + messages := al.contextBuilder.BuildMessages( - al.sessions.GetHistory(msg.SessionKey), + history, + summary, msg.Content, nil, ) @@ -187,7 +195,73 @@ func (al *AgentLoop) processMessage(ctx context.Context, msg bus.InboundMessage) al.sessions.AddMessage(msg.SessionKey, "user", msg.Content) al.sessions.AddMessage(msg.SessionKey, "assistant", finalContent) + + // Context compression logic + newHistory := al.sessions.GetHistory(msg.SessionKey) + if len(newHistory) > 20 { + if _, loading := al.summarizing.LoadOrStore(msg.SessionKey, true); !loading { + go func() { + defer al.summarizing.Delete(msg.SessionKey) + al.summarizeSession(msg.SessionKey) + }() + } + } + al.sessions.Save(al.sessions.GetOrCreate(msg.SessionKey)) return finalContent, nil } + +func (al *AgentLoop) summarizeSession(sessionKey string) { + ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) + defer cancel() + + history := al.sessions.GetHistory(sessionKey) + summary := al.sessions.GetSummary(sessionKey) + + // Keep last 4 messages, summarize the rest + if len(history) <= 4 { + return + } + + toSummarize := history[:len(history)-4] + + prompt := "Below is a conversation history and an optional existing summary. " + + "Please provide a concise summary of the conversation so far, " + + "preserving the core context and key points discussed. " + + "If there's an existing summary, incorporate it into the new one.\n\n" + + if summary != "" { + prompt += "EXISTING SUMMARY: " + summary + "\n\n" + } + + prompt += "CONVERSATION TO SUMMARIZE:\n" + for _, m := range toSummarize { + if m.Role == "user" || m.Role == "assistant" { + prompt += fmt.Sprintf("%s: %s\n", m.Role, m.Content) + } + } + + messages := []providers.Message{ + { + Role: "user", + Content: prompt, + }, + } + + response, err := al.provider.Chat(ctx, messages, nil, al.model, map[string]interface{}{ + "max_tokens": 1024, + "temperature": 0.3, + }) + + if err != nil { + fmt.Printf("Error summarizing session %s: %v\n", sessionKey, err) + return + } + + if response.Content != "" { + al.sessions.SetSummary(sessionKey, response.Content) + al.sessions.TruncateHistory(sessionKey, 4) + al.sessions.Save(al.sessions.GetOrCreate(sessionKey)) + } +} diff --git a/pkg/session/manager.go b/pkg/session/manager.go index 9e17b30..df86724 100644 --- a/pkg/session/manager.go +++ b/pkg/session/manager.go @@ -13,6 +13,7 @@ import ( type Session struct { Key string `json:"key"` Messages []providers.Message `json:"messages"` + Summary string `json:"summary,omitempty"` Created time.Time `json:"created"` Updated time.Time `json:"updated"` } @@ -92,6 +93,45 @@ func (sm *SessionManager) GetHistory(key string) []providers.Message { return history } +func (sm *SessionManager) GetSummary(key string) string { + sm.mu.RLock() + defer sm.mu.RUnlock() + + session, ok := sm.sessions[key] + if !ok { + return "" + } + return session.Summary +} + +func (sm *SessionManager) SetSummary(key string, summary string) { + sm.mu.Lock() + defer sm.mu.Unlock() + + session, ok := sm.sessions[key] + if ok { + session.Summary = summary + session.Updated = time.Now() + } +} + +func (sm *SessionManager) TruncateHistory(key string, keepLast int) { + sm.mu.Lock() + defer sm.mu.Unlock() + + session, ok := sm.sessions[key] + if !ok { + return + } + + if len(session.Messages) <= keepLast { + return + } + + session.Messages = session.Messages[len(session.Messages)-keepLast:] + session.Updated = time.Now() +} + func (sm *SessionManager) Save(session *Session) error { if sm.storage == "" { return nil