feat(memory): add observational memory system
Implement an observation pipeline that passively monitors agent interactions to extract, score, and persist notable events. The reflector periodically consolidates observations into higher-order insights stored in archival memory.
This commit is contained in:
parent
b5ab7a6108
commit
542b78a49d
7 changed files with 922 additions and 0 deletions
136
pkg/memory/observation/manager.go
Normal file
136
pkg/memory/observation/manager.go
Normal file
|
|
@ -0,0 +1,136 @@
|
||||||
|
package observation
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"log/slog"
|
||||||
|
"sync"
|
||||||
|
|
||||||
|
"github.com/sipeed/picoclaw/pkg/memory"
|
||||||
|
)
|
||||||
|
|
||||||
|
// Manager orchestrates the observation lifecycle: threshold detection,
|
||||||
|
// observer/reflector invocation, persistence, and async execution.
|
||||||
|
type Manager struct {
|
||||||
|
store *Store
|
||||||
|
observer *Observer
|
||||||
|
reflector *Reflector
|
||||||
|
|
||||||
|
mu sync.Mutex
|
||||||
|
running map[string]bool // sessionKey -> running flag to prevent concurrent runs
|
||||||
|
}
|
||||||
|
|
||||||
|
// ManagerConfig bundles configuration for the observation system.
|
||||||
|
type ManagerConfig struct {
|
||||||
|
Observer ObserverConfig
|
||||||
|
Reflector ReflectorConfig
|
||||||
|
}
|
||||||
|
|
||||||
|
func DefaultManagerConfig() ManagerConfig {
|
||||||
|
return ManagerConfig{
|
||||||
|
Observer: DefaultObserverConfig(),
|
||||||
|
Reflector: DefaultReflectorConfig(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewManager creates an observation Manager backed by the given delegate.
|
||||||
|
func NewManager(delegate memory.MemoryDelegate, agentID string, callModel ModelFunc, cfg ManagerConfig) *Manager {
|
||||||
|
store := NewStore(delegate, agentID)
|
||||||
|
return &Manager{
|
||||||
|
store: store,
|
||||||
|
observer: NewObserver(callModel, cfg.Observer),
|
||||||
|
reflector: NewReflector(callModel, cfg.Reflector),
|
||||||
|
running: make(map[string]bool),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// MaybeObserveAsync checks if observation is needed and runs it in a background
|
||||||
|
// goroutine if so. Non-blocking; safe to call on every agent turn.
|
||||||
|
func (m *Manager) MaybeObserveAsync(ctx context.Context, sessionKey string, tailMessages []MessagePair) {
|
||||||
|
if !m.observer.ShouldObserve(tailMessages) {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
if !m.tryAcquire(sessionKey) {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
// Copy messages to avoid data races with the caller.
|
||||||
|
msgs := make([]MessagePair, len(tailMessages))
|
||||||
|
copy(msgs, tailMessages)
|
||||||
|
|
||||||
|
go func() {
|
||||||
|
defer m.release(sessionKey)
|
||||||
|
m.runObservation(ctx, sessionKey, msgs)
|
||||||
|
}()
|
||||||
|
}
|
||||||
|
|
||||||
|
// LoadBlock returns the formatted observation block for system prompt injection.
|
||||||
|
// Returns empty string if no observations exist.
|
||||||
|
func (m *Manager) LoadBlock(ctx context.Context, sessionKey string) string {
|
||||||
|
obs, err := m.store.Load(ctx, sessionKey)
|
||||||
|
if err != nil {
|
||||||
|
slog.Warn("failed to load observations", "session", sessionKey, "error", err)
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
return FormatBlock(obs)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Store returns the underlying observation store for direct access.
|
||||||
|
func (m *Manager) Store() *Store {
|
||||||
|
return m.store
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m *Manager) runObservation(ctx context.Context, sessionKey string, messages []MessagePair) {
|
||||||
|
existing, err := m.store.Load(ctx, sessionKey)
|
||||||
|
if err != nil {
|
||||||
|
slog.Error("observation: failed to load existing", "session", sessionKey, "error", err)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
newObs, err := m.observer.Observe(ctx, messages, existing)
|
||||||
|
if err != nil {
|
||||||
|
slog.Error("observation: observer failed", "session", sessionKey, "error", err)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
if len(newObs) == 0 {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
all := append(existing, newObs...)
|
||||||
|
|
||||||
|
if m.reflector.ShouldReflect(all) {
|
||||||
|
pruned, err := m.reflector.Reflect(ctx, all)
|
||||||
|
if err != nil {
|
||||||
|
slog.Error("observation: reflector failed", "session", sessionKey, "error", err)
|
||||||
|
// Save unpruned observations rather than losing them
|
||||||
|
} else {
|
||||||
|
all = pruned
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := m.store.Save(ctx, sessionKey, all); err != nil {
|
||||||
|
slog.Error("observation: failed to save", "session", sessionKey, "error", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
slog.Info("observation: updated",
|
||||||
|
"session", sessionKey,
|
||||||
|
"new", len(newObs),
|
||||||
|
"total", len(all))
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m *Manager) tryAcquire(sessionKey string) bool {
|
||||||
|
m.mu.Lock()
|
||||||
|
defer m.mu.Unlock()
|
||||||
|
if m.running[sessionKey] {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
m.running[sessionKey] = true
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m *Manager) release(sessionKey string) {
|
||||||
|
m.mu.Lock()
|
||||||
|
defer m.mu.Unlock()
|
||||||
|
delete(m.running, sessionKey)
|
||||||
|
}
|
||||||
150
pkg/memory/observation/observation.go
Normal file
150
pkg/memory/observation/observation.go
Normal file
|
|
@ -0,0 +1,150 @@
|
||||||
|
// Package observation implements Mastra-style Observational Memory.
|
||||||
|
//
|
||||||
|
// Raw conversation is compressed into prioritized observations that form a
|
||||||
|
// stable, prompt-cacheable prefix in the system prompt. Two background agents
|
||||||
|
// maintain the observation block:
|
||||||
|
//
|
||||||
|
// - Observer: fires when the uncompressed tail exceeds a token threshold,
|
||||||
|
// compressing recent messages into new observations.
|
||||||
|
// - Reflector: fires when the observation block itself exceeds a threshold,
|
||||||
|
// garbage-collecting low-priority observations.
|
||||||
|
//
|
||||||
|
// Each observation carries a three-date model: observation date (when created),
|
||||||
|
// referenced date (when the event occurred), and a human-readable relative date.
|
||||||
|
package observation
|
||||||
|
|
||||||
|
import (
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
// Priority encodes the importance of an observation for prompt display.
|
||||||
|
type Priority string
|
||||||
|
|
||||||
|
const (
|
||||||
|
PriorityCritical Priority = "critical" // 🔴
|
||||||
|
PriorityNotable Priority = "notable" // 🟡
|
||||||
|
PriorityInformational Priority = "informational" // 🔵
|
||||||
|
)
|
||||||
|
|
||||||
|
func (p Priority) Emoji() string {
|
||||||
|
switch p {
|
||||||
|
case PriorityCritical:
|
||||||
|
return "🔴"
|
||||||
|
case PriorityNotable:
|
||||||
|
return "🟡"
|
||||||
|
case PriorityInformational:
|
||||||
|
return "🔵"
|
||||||
|
default:
|
||||||
|
return "🔵"
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Observation is a single compressed insight extracted from conversation.
|
||||||
|
type Observation struct {
|
||||||
|
Content string `json:"content"`
|
||||||
|
Priority Priority `json:"priority"`
|
||||||
|
ObservedAt int64 `json:"observed_at"` // Unix timestamp: when observation was created
|
||||||
|
ReferencedAt int64 `json:"referenced_at"` // Unix timestamp: when the referenced event occurred
|
||||||
|
RelativeDate string `json:"relative_date"` // Human-readable: "2 days ago", "today", etc.
|
||||||
|
}
|
||||||
|
|
||||||
|
// ObservedTime returns ObservedAt as time.Time.
|
||||||
|
func (o Observation) ObservedTime() time.Time {
|
||||||
|
return time.Unix(o.ObservedAt, 0)
|
||||||
|
}
|
||||||
|
|
||||||
|
// ReferencedTime returns ReferencedAt as time.Time.
|
||||||
|
func (o Observation) ReferencedTime() time.Time {
|
||||||
|
return time.Unix(o.ReferencedAt, 0)
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewObservation creates an observation with the three-date model.
|
||||||
|
// referenceTime is when the observed event happened; observeTime is now.
|
||||||
|
func NewObservation(content string, priority Priority, referenceTime, observeTime time.Time) Observation {
|
||||||
|
return Observation{
|
||||||
|
Content: content,
|
||||||
|
Priority: priority,
|
||||||
|
ObservedAt: observeTime.Unix(),
|
||||||
|
ReferencedAt: referenceTime.Unix(),
|
||||||
|
RelativeDate: relativeDate(referenceTime, observeTime),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// FormatBlock renders the entire observation list as a single text block
|
||||||
|
// suitable for system prompt injection. Format per line:
|
||||||
|
//
|
||||||
|
// DATE EMOJI HH:MM observation_text
|
||||||
|
func FormatBlock(observations []Observation) string {
|
||||||
|
if len(observations) == 0 {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
|
||||||
|
var sb strings.Builder
|
||||||
|
for _, o := range observations {
|
||||||
|
ref := o.ReferencedTime()
|
||||||
|
sb.WriteString(fmt.Sprintf("%s %s %s %s\n",
|
||||||
|
ref.Format("2006-01-02"),
|
||||||
|
o.Priority.Emoji(),
|
||||||
|
ref.Format("15:04"),
|
||||||
|
o.Content,
|
||||||
|
))
|
||||||
|
}
|
||||||
|
return sb.String()
|
||||||
|
}
|
||||||
|
|
||||||
|
// MarshalObservations serializes observations to JSON for KV storage.
|
||||||
|
func MarshalObservations(obs []Observation) (string, error) {
|
||||||
|
data, err := json.Marshal(obs)
|
||||||
|
if err != nil {
|
||||||
|
return "", fmt.Errorf("marshal observations: %w", err)
|
||||||
|
}
|
||||||
|
return string(data), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// UnmarshalObservations deserializes observations from KV storage.
|
||||||
|
func UnmarshalObservations(data string) ([]Observation, error) {
|
||||||
|
if data == "" {
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
|
var obs []Observation
|
||||||
|
if err := json.Unmarshal([]byte(data), &obs); err != nil {
|
||||||
|
return nil, fmt.Errorf("unmarshal observations: %w", err)
|
||||||
|
}
|
||||||
|
return obs, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func relativeDate(ref, now time.Time) string {
|
||||||
|
diff := now.Sub(ref)
|
||||||
|
switch {
|
||||||
|
case diff < time.Minute:
|
||||||
|
return "just now"
|
||||||
|
case diff < time.Hour:
|
||||||
|
m := int(diff.Minutes())
|
||||||
|
if m == 1 {
|
||||||
|
return "1 minute ago"
|
||||||
|
}
|
||||||
|
return fmt.Sprintf("%d minutes ago", m)
|
||||||
|
case diff < 24*time.Hour:
|
||||||
|
h := int(diff.Hours())
|
||||||
|
if h == 1 {
|
||||||
|
return "1 hour ago"
|
||||||
|
}
|
||||||
|
return fmt.Sprintf("%d hours ago", h)
|
||||||
|
case diff < 48*time.Hour:
|
||||||
|
return "yesterday"
|
||||||
|
case diff < 7*24*time.Hour:
|
||||||
|
d := int(diff.Hours() / 24)
|
||||||
|
return fmt.Sprintf("%d days ago", d)
|
||||||
|
case diff < 30*24*time.Hour:
|
||||||
|
w := int(diff.Hours() / (24 * 7))
|
||||||
|
if w == 1 {
|
||||||
|
return "1 week ago"
|
||||||
|
}
|
||||||
|
return fmt.Sprintf("%d weeks ago", w)
|
||||||
|
default:
|
||||||
|
return ref.Format("2006-01-02")
|
||||||
|
}
|
||||||
|
}
|
||||||
251
pkg/memory/observation/observation_test.go
Normal file
251
pkg/memory/observation/observation_test.go
Normal file
|
|
@ -0,0 +1,251 @@
|
||||||
|
package observation
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"fmt"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/stretchr/testify/assert"
|
||||||
|
"github.com/stretchr/testify/require"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestNewObservation_ThreeDateModel(t *testing.T) {
|
||||||
|
ref := time.Date(2026, 2, 16, 10, 0, 0, 0, time.UTC)
|
||||||
|
obs := time.Date(2026, 2, 18, 14, 30, 0, 0, time.UTC)
|
||||||
|
|
||||||
|
o := NewObservation("User prefers Go over Rust", PriorityNotable, ref, obs)
|
||||||
|
|
||||||
|
assert.Equal(t, "User prefers Go over Rust", o.Content)
|
||||||
|
assert.Equal(t, PriorityNotable, o.Priority)
|
||||||
|
assert.Equal(t, ref.Unix(), o.ReferencedAt)
|
||||||
|
assert.Equal(t, obs.Unix(), o.ObservedAt)
|
||||||
|
assert.Equal(t, "2 days ago", o.RelativeDate)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRelativeDate(t *testing.T) {
|
||||||
|
now := time.Date(2026, 2, 18, 14, 0, 0, 0, time.UTC)
|
||||||
|
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
ref time.Time
|
||||||
|
want string
|
||||||
|
}{
|
||||||
|
{"just now", now.Add(-30 * time.Second), "just now"},
|
||||||
|
{"1 minute ago", now.Add(-1 * time.Minute), "1 minute ago"},
|
||||||
|
{"5 minutes ago", now.Add(-5 * time.Minute), "5 minutes ago"},
|
||||||
|
{"1 hour ago", now.Add(-1 * time.Hour), "1 hour ago"},
|
||||||
|
{"3 hours ago", now.Add(-3 * time.Hour), "3 hours ago"},
|
||||||
|
{"yesterday", now.Add(-30 * time.Hour), "yesterday"},
|
||||||
|
{"3 days ago", now.Add(-3 * 24 * time.Hour), "3 days ago"},
|
||||||
|
{"1 week ago", now.Add(-7 * 24 * time.Hour), "1 week ago"},
|
||||||
|
{"3 weeks ago", now.Add(-21 * 24 * time.Hour), "3 weeks ago"},
|
||||||
|
{"old date", now.Add(-60 * 24 * time.Hour), "2025-12-20"},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tc := range tests {
|
||||||
|
t.Run(tc.name, func(t *testing.T) {
|
||||||
|
got := relativeDate(tc.ref, now)
|
||||||
|
assert.Equal(t, tc.want, got)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestPriorityEmoji(t *testing.T) {
|
||||||
|
assert.Equal(t, "🔴", PriorityCritical.Emoji())
|
||||||
|
assert.Equal(t, "🟡", PriorityNotable.Emoji())
|
||||||
|
assert.Equal(t, "🔵", PriorityInformational.Emoji())
|
||||||
|
assert.Equal(t, "🔵", Priority("unknown").Emoji())
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFormatBlock(t *testing.T) {
|
||||||
|
now := time.Date(2026, 2, 18, 14, 30, 0, 0, time.UTC)
|
||||||
|
obs := []Observation{
|
||||||
|
NewObservation("Decision: use SQLite", PriorityCritical, now, now),
|
||||||
|
NewObservation("Prefers Go", PriorityNotable, now.Add(-time.Hour), now),
|
||||||
|
}
|
||||||
|
|
||||||
|
block := FormatBlock(obs)
|
||||||
|
assert.Contains(t, block, "🔴")
|
||||||
|
assert.Contains(t, block, "🟡")
|
||||||
|
assert.Contains(t, block, "Decision: use SQLite")
|
||||||
|
assert.Contains(t, block, "Prefers Go")
|
||||||
|
assert.Contains(t, block, "2026-02-18")
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFormatBlock_Empty(t *testing.T) {
|
||||||
|
assert.Equal(t, "", FormatBlock(nil))
|
||||||
|
assert.Equal(t, "", FormatBlock([]Observation{}))
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestMarshalUnmarshalRoundTrip(t *testing.T) {
|
||||||
|
now := time.Now()
|
||||||
|
obs := []Observation{
|
||||||
|
NewObservation("Fact A", PriorityCritical, now, now),
|
||||||
|
NewObservation("Fact B", PriorityInformational, now.Add(-time.Hour), now),
|
||||||
|
}
|
||||||
|
|
||||||
|
data, err := MarshalObservations(obs)
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.NotEmpty(t, data)
|
||||||
|
|
||||||
|
parsed, err := UnmarshalObservations(data)
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Len(t, parsed, 2)
|
||||||
|
assert.Equal(t, "Fact A", parsed[0].Content)
|
||||||
|
assert.Equal(t, PriorityCritical, parsed[0].Priority)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestUnmarshalObservations_Empty(t *testing.T) {
|
||||||
|
obs, err := UnmarshalObservations("")
|
||||||
|
assert.NoError(t, err)
|
||||||
|
assert.Nil(t, obs)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestEstimateTokens(t *testing.T) {
|
||||||
|
tokens := EstimateTokens("Hello, world!")
|
||||||
|
assert.True(t, tokens > 0)
|
||||||
|
assert.True(t, tokens < 20)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestParseObservations(t *testing.T) {
|
||||||
|
now := time.Now()
|
||||||
|
response := `critical|User decided to migrate to SQLite
|
||||||
|
notable|Prefers hexagonal architecture
|
||||||
|
informational|Uses VS Code as primary editor
|
||||||
|
invalid line without pipe
|
||||||
|
notable|`
|
||||||
|
|
||||||
|
obs := parseObservations(response, now)
|
||||||
|
assert.Len(t, obs, 3)
|
||||||
|
assert.Equal(t, PriorityCritical, obs[0].Priority)
|
||||||
|
assert.Equal(t, "User decided to migrate to SQLite", obs[0].Content)
|
||||||
|
assert.Equal(t, PriorityNotable, obs[1].Priority)
|
||||||
|
assert.Equal(t, PriorityInformational, obs[2].Priority)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestObserver_ShouldObserve(t *testing.T) {
|
||||||
|
mockModel := func(_ context.Context, _ string) (string, error) {
|
||||||
|
return "", nil
|
||||||
|
}
|
||||||
|
|
||||||
|
o := NewObserver(mockModel, ObserverConfig{TokenThreshold: 100})
|
||||||
|
|
||||||
|
small := []MessagePair{{Role: "user", Content: "Hi"}}
|
||||||
|
assert.False(t, o.ShouldObserve(small))
|
||||||
|
|
||||||
|
large := []MessagePair{{Role: "user", Content: strings.Repeat("word ", 200)}}
|
||||||
|
assert.True(t, o.ShouldObserve(large))
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestObserver_Observe(t *testing.T) {
|
||||||
|
mockModel := func(_ context.Context, prompt string) (string, error) {
|
||||||
|
return "critical|Important decision made\nnotable|User preference noted", nil
|
||||||
|
}
|
||||||
|
|
||||||
|
o := NewObserver(mockModel, DefaultObserverConfig())
|
||||||
|
|
||||||
|
msgs := []MessagePair{
|
||||||
|
{Role: "user", Content: "I want to use SQLite for everything"},
|
||||||
|
{Role: "assistant", Content: "Good choice for embedded use cases"},
|
||||||
|
}
|
||||||
|
|
||||||
|
obs, err := o.Observe(context.Background(), msgs, nil)
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Len(t, obs, 2)
|
||||||
|
assert.Equal(t, PriorityCritical, obs[0].Priority)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestReflector_ShouldReflect(t *testing.T) {
|
||||||
|
mockModel := func(_ context.Context, _ string) (string, error) {
|
||||||
|
return "", nil
|
||||||
|
}
|
||||||
|
|
||||||
|
r := NewReflector(mockModel, ReflectorConfig{TokenThreshold: 100})
|
||||||
|
|
||||||
|
small := []Observation{NewObservation("Small fact", PriorityInformational, time.Now(), time.Now())}
|
||||||
|
assert.False(t, r.ShouldReflect(small))
|
||||||
|
|
||||||
|
var large []Observation
|
||||||
|
for i := 0; i < 50; i++ {
|
||||||
|
large = append(large, NewObservation(strings.Repeat("word ", 20), PriorityInformational, time.Now(), time.Now()))
|
||||||
|
}
|
||||||
|
assert.True(t, r.ShouldReflect(large))
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestReflector_Reflect(t *testing.T) {
|
||||||
|
mockModel := func(_ context.Context, prompt string) (string, error) {
|
||||||
|
return "KEEP 0\nDROP 1\nKEEP 2", nil
|
||||||
|
}
|
||||||
|
|
||||||
|
r := NewReflector(mockModel, DefaultReflectorConfig())
|
||||||
|
now := time.Now()
|
||||||
|
|
||||||
|
obs := []Observation{
|
||||||
|
NewObservation("Critical fact", PriorityCritical, now, now),
|
||||||
|
NewObservation("Old info", PriorityInformational, now.Add(-24*time.Hour), now),
|
||||||
|
NewObservation("Notable thing", PriorityNotable, now, now),
|
||||||
|
}
|
||||||
|
|
||||||
|
kept, err := r.Reflect(context.Background(), obs)
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Len(t, kept, 2)
|
||||||
|
assert.Equal(t, "Critical fact", kept[0].Content)
|
||||||
|
assert.Equal(t, "Notable thing", kept[1].Content)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestReflector_ReflectFallbackKeepsCritical(t *testing.T) {
|
||||||
|
mockModel := func(_ context.Context, _ string) (string, error) {
|
||||||
|
return "garbage output", nil
|
||||||
|
}
|
||||||
|
|
||||||
|
r := NewReflector(mockModel, DefaultReflectorConfig())
|
||||||
|
now := time.Now()
|
||||||
|
|
||||||
|
obs := []Observation{
|
||||||
|
NewObservation("Must keep", PriorityCritical, now, now),
|
||||||
|
NewObservation("Can drop", PriorityInformational, now, now),
|
||||||
|
}
|
||||||
|
|
||||||
|
kept, err := r.Reflect(context.Background(), obs)
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Len(t, kept, 1)
|
||||||
|
assert.Equal(t, "Must keep", kept[0].Content)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestParsePriority(t *testing.T) {
|
||||||
|
assert.Equal(t, PriorityCritical, parsePriority("critical"))
|
||||||
|
assert.Equal(t, PriorityCritical, parsePriority("CRITICAL"))
|
||||||
|
assert.Equal(t, PriorityNotable, parsePriority("notable"))
|
||||||
|
assert.Equal(t, PriorityInformational, parsePriority("informational"))
|
||||||
|
assert.Equal(t, PriorityInformational, parsePriority("unknown"))
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestParseKeptIndices(t *testing.T) {
|
||||||
|
now := time.Now()
|
||||||
|
obs := []Observation{
|
||||||
|
NewObservation("A", PriorityCritical, now, now),
|
||||||
|
NewObservation("B", PriorityNotable, now, now),
|
||||||
|
NewObservation("C", PriorityInformational, now, now),
|
||||||
|
}
|
||||||
|
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
response string
|
||||||
|
wantLen int
|
||||||
|
}{
|
||||||
|
{"normal", "KEEP 0\nDROP 1\nKEEP 2", 2},
|
||||||
|
{"all keep", "KEEP 0\nKEEP 1\nKEEP 2", 3},
|
||||||
|
{"all drop", "DROP 0\nDROP 1\nDROP 2", 1}, // Fallback keeps critical
|
||||||
|
{"invalid output", "blah blah", 1}, // Fallback keeps critical
|
||||||
|
{"out of range", "KEEP 99", 1}, // Fallback keeps critical
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tc := range tests {
|
||||||
|
t.Run(tc.name, func(t *testing.T) {
|
||||||
|
result := parseKeptIndices(tc.response, obs)
|
||||||
|
assert.Len(t, result, tc.wantLen, fmt.Sprintf("response: %q", tc.response))
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
140
pkg/memory/observation/observer.go
Normal file
140
pkg/memory/observation/observer.go
Normal file
|
|
@ -0,0 +1,140 @@
|
||||||
|
package observation
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"fmt"
|
||||||
|
"log/slog"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
// ModelFunc is a function that sends a prompt to an LLM and returns the text response.
|
||||||
|
// This decouples the observation package from the fantasy/model layer.
|
||||||
|
type ModelFunc func(ctx context.Context, prompt string) (string, error)
|
||||||
|
|
||||||
|
// ObserverConfig controls when and how the Observer triggers.
|
||||||
|
type ObserverConfig struct {
|
||||||
|
TokenThreshold int // Uncompressed tail token count to trigger observation (default 30000)
|
||||||
|
}
|
||||||
|
|
||||||
|
func DefaultObserverConfig() ObserverConfig {
|
||||||
|
return ObserverConfig{
|
||||||
|
TokenThreshold: 30000,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Observer compresses raw conversation into prioritized observations.
|
||||||
|
// It fires when the uncompressed tail exceeds the token threshold.
|
||||||
|
type Observer struct {
|
||||||
|
callModel ModelFunc
|
||||||
|
cfg ObserverConfig
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewObserver(callModel ModelFunc, cfg ObserverConfig) *Observer {
|
||||||
|
return &Observer{
|
||||||
|
callModel: callModel,
|
||||||
|
cfg: cfg,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ShouldObserve returns true if the uncompressed tail has exceeded the token threshold.
|
||||||
|
func (o *Observer) ShouldObserve(tailMessages []MessagePair) bool {
|
||||||
|
return EstimateMessagesTokens(tailMessages) >= o.cfg.TokenThreshold
|
||||||
|
}
|
||||||
|
|
||||||
|
// Observe compresses the given messages into a list of new observations.
|
||||||
|
// The LLM is asked to extract key insights, decisions, and facts.
|
||||||
|
func (o *Observer) Observe(ctx context.Context, messages []MessagePair, existingObs []Observation) ([]Observation, error) {
|
||||||
|
if len(messages) == 0 {
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
prompt := buildObserverPrompt(messages, existingObs)
|
||||||
|
|
||||||
|
resp, err := o.callModel(ctx, prompt)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("observer LLM call: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
return parseObservations(resp, time.Now()), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func buildObserverPrompt(messages []MessagePair, existing []Observation) string {
|
||||||
|
var sb strings.Builder
|
||||||
|
|
||||||
|
sb.WriteString(`You are an observation extractor. Analyze the conversation below and extract key observations.
|
||||||
|
|
||||||
|
For each observation, output one line in this exact format:
|
||||||
|
PRIORITY|CONTENT
|
||||||
|
|
||||||
|
Where PRIORITY is one of: critical, notable, informational
|
||||||
|
|
||||||
|
Rules:
|
||||||
|
- Extract 3-10 observations from the conversation
|
||||||
|
- Critical: decisions made, errors encountered, important commitments
|
||||||
|
- Notable: useful information learned, preferences expressed, patterns identified
|
||||||
|
- Informational: context details, minor facts, status updates
|
||||||
|
- Be concise: each observation should be 1-2 sentences max
|
||||||
|
- Focus on facts and insights, not conversation flow
|
||||||
|
- Do NOT include observations that duplicate existing ones
|
||||||
|
|
||||||
|
`)
|
||||||
|
|
||||||
|
if len(existing) > 0 {
|
||||||
|
sb.WriteString("## Existing observations (do not duplicate):\n")
|
||||||
|
for _, o := range existing {
|
||||||
|
sb.WriteString(fmt.Sprintf("- %s\n", o.Content))
|
||||||
|
}
|
||||||
|
sb.WriteString("\n")
|
||||||
|
}
|
||||||
|
|
||||||
|
sb.WriteString("## Conversation to analyze:\n\n")
|
||||||
|
for _, m := range messages {
|
||||||
|
sb.WriteString(fmt.Sprintf("[%s]: %s\n\n", m.Role, m.Content))
|
||||||
|
}
|
||||||
|
|
||||||
|
return sb.String()
|
||||||
|
}
|
||||||
|
|
||||||
|
func parseObservations(response string, now time.Time) []Observation {
|
||||||
|
var observations []Observation
|
||||||
|
|
||||||
|
for _, line := range strings.Split(response, "\n") {
|
||||||
|
line = strings.TrimSpace(line)
|
||||||
|
if line == "" {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
parts := strings.SplitN(line, "|", 2)
|
||||||
|
if len(parts) != 2 {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
priority := parsePriority(strings.TrimSpace(parts[0]))
|
||||||
|
content := strings.TrimSpace(parts[1])
|
||||||
|
if content == "" {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
observations = append(observations, NewObservation(content, priority, now, now))
|
||||||
|
}
|
||||||
|
|
||||||
|
if len(observations) == 0 {
|
||||||
|
slog.Warn("observer produced no parseable observations from LLM response")
|
||||||
|
}
|
||||||
|
|
||||||
|
return observations
|
||||||
|
}
|
||||||
|
|
||||||
|
func parsePriority(s string) Priority {
|
||||||
|
switch strings.ToLower(s) {
|
||||||
|
case "critical":
|
||||||
|
return PriorityCritical
|
||||||
|
case "notable":
|
||||||
|
return PriorityNotable
|
||||||
|
case "informational":
|
||||||
|
return PriorityInformational
|
||||||
|
default:
|
||||||
|
return PriorityInformational
|
||||||
|
}
|
||||||
|
}
|
||||||
138
pkg/memory/observation/reflector.go
Normal file
138
pkg/memory/observation/reflector.go
Normal file
|
|
@ -0,0 +1,138 @@
|
||||||
|
package observation
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"fmt"
|
||||||
|
"log/slog"
|
||||||
|
"strings"
|
||||||
|
)
|
||||||
|
|
||||||
|
// ReflectorConfig controls when and how the Reflector triggers.
|
||||||
|
type ReflectorConfig struct {
|
||||||
|
TokenThreshold int // Observation block token count to trigger reflection (default 40000)
|
||||||
|
}
|
||||||
|
|
||||||
|
func DefaultReflectorConfig() ReflectorConfig {
|
||||||
|
return ReflectorConfig{
|
||||||
|
TokenThreshold: 40000,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Reflector garbage-collects low-priority observations when the block exceeds
|
||||||
|
// its token threshold. This is the only operation that invalidates the full
|
||||||
|
// prompt cache (rare by design).
|
||||||
|
type Reflector struct {
|
||||||
|
callModel ModelFunc
|
||||||
|
cfg ReflectorConfig
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewReflector(callModel ModelFunc, cfg ReflectorConfig) *Reflector {
|
||||||
|
return &Reflector{
|
||||||
|
callModel: callModel,
|
||||||
|
cfg: cfg,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ShouldReflect returns true if the observation block exceeds the threshold.
|
||||||
|
func (r *Reflector) ShouldReflect(observations []Observation) bool {
|
||||||
|
block := FormatBlock(observations)
|
||||||
|
return EstimateTokens(block) >= r.cfg.TokenThreshold
|
||||||
|
}
|
||||||
|
|
||||||
|
// Reflect asks the LLM to select observations worth keeping.
|
||||||
|
// Returns the pruned observation list.
|
||||||
|
func (r *Reflector) Reflect(ctx context.Context, observations []Observation) ([]Observation, error) {
|
||||||
|
if len(observations) == 0 {
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
prompt := buildReflectorPrompt(observations)
|
||||||
|
|
||||||
|
resp, err := r.callModel(ctx, prompt)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("reflector LLM call: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
kept := parseKeptIndices(resp, observations)
|
||||||
|
|
||||||
|
slog.Info("reflector GC",
|
||||||
|
"before", len(observations),
|
||||||
|
"after", len(kept),
|
||||||
|
"dropped", len(observations)-len(kept))
|
||||||
|
|
||||||
|
return kept, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func buildReflectorPrompt(observations []Observation) string {
|
||||||
|
var sb strings.Builder
|
||||||
|
|
||||||
|
sb.WriteString(`You are a memory curator. Review the observations below and decide which to KEEP.
|
||||||
|
|
||||||
|
Rules:
|
||||||
|
- KEEP all critical observations
|
||||||
|
- KEEP notable observations that are still relevant
|
||||||
|
- DROP informational observations that are outdated or redundant
|
||||||
|
- DROP observations that have been superseded by newer ones
|
||||||
|
- Aim to reduce the list by 30-50%
|
||||||
|
|
||||||
|
For each observation, output KEEP or DROP followed by the index number:
|
||||||
|
KEEP 0
|
||||||
|
DROP 1
|
||||||
|
KEEP 2
|
||||||
|
...
|
||||||
|
|
||||||
|
## Observations:
|
||||||
|
|
||||||
|
`)
|
||||||
|
|
||||||
|
for i, o := range observations {
|
||||||
|
sb.WriteString(fmt.Sprintf("[%d] %s %s — %s (%s)\n",
|
||||||
|
i,
|
||||||
|
o.Priority.Emoji(),
|
||||||
|
o.Priority,
|
||||||
|
o.Content,
|
||||||
|
o.RelativeDate,
|
||||||
|
))
|
||||||
|
}
|
||||||
|
|
||||||
|
return sb.String()
|
||||||
|
}
|
||||||
|
|
||||||
|
func parseKeptIndices(response string, observations []Observation) []Observation {
|
||||||
|
kept := make(map[int]bool)
|
||||||
|
|
||||||
|
for _, line := range strings.Split(response, "\n") {
|
||||||
|
line = strings.TrimSpace(line)
|
||||||
|
if line == "" {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
var action string
|
||||||
|
var idx int
|
||||||
|
if _, err := fmt.Sscanf(line, "%s %d", &action, &idx); err != nil {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
if strings.EqualFold(action, "KEEP") && idx >= 0 && idx < len(observations) {
|
||||||
|
kept[idx] = true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// If parsing failed or nothing was kept, keep all critical observations at minimum
|
||||||
|
if len(kept) == 0 {
|
||||||
|
slog.Warn("reflector produced no valid KEEP instructions, preserving all critical observations")
|
||||||
|
for i, o := range observations {
|
||||||
|
if o.Priority == PriorityCritical {
|
||||||
|
kept[i] = true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
var result []Observation
|
||||||
|
for i, o := range observations {
|
||||||
|
if kept[i] {
|
||||||
|
result = append(result, o)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return result
|
||||||
|
}
|
||||||
57
pkg/memory/observation/store.go
Normal file
57
pkg/memory/observation/store.go
Normal file
|
|
@ -0,0 +1,57 @@
|
||||||
|
package observation
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"fmt"
|
||||||
|
|
||||||
|
"github.com/sipeed/picoclaw/pkg/memory"
|
||||||
|
)
|
||||||
|
|
||||||
|
const kvPrefix = "obs:"
|
||||||
|
|
||||||
|
// Store persists observations to the agent_kv table via MemoryDelegate.
|
||||||
|
// Each session's observations are stored as a single JSON array under
|
||||||
|
// the key "obs:<session_key>".
|
||||||
|
type Store struct {
|
||||||
|
delegate memory.MemoryDelegate
|
||||||
|
agentID string
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewStore(delegate memory.MemoryDelegate, agentID string) *Store {
|
||||||
|
return &Store{
|
||||||
|
delegate: delegate,
|
||||||
|
agentID: agentID,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *Store) key(sessionKey string) string {
|
||||||
|
return kvPrefix + sessionKey
|
||||||
|
}
|
||||||
|
|
||||||
|
// Load retrieves all observations for a session.
|
||||||
|
func (s *Store) Load(ctx context.Context, sessionKey string) ([]Observation, error) {
|
||||||
|
data, err := s.delegate.GetKV(ctx, s.agentID, s.key(sessionKey))
|
||||||
|
if err != nil {
|
||||||
|
return nil, nil // Key not found is not an error
|
||||||
|
}
|
||||||
|
return UnmarshalObservations(data)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Save persists the full observation list for a session.
|
||||||
|
func (s *Store) Save(ctx context.Context, sessionKey string, obs []Observation) error {
|
||||||
|
data, err := MarshalObservations(obs)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("save observations: %w", err)
|
||||||
|
}
|
||||||
|
return s.delegate.UpsertKV(ctx, s.agentID, s.key(sessionKey), data)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Append adds new observations to the existing list and persists.
|
||||||
|
func (s *Store) Append(ctx context.Context, sessionKey string, newObs []Observation) error {
|
||||||
|
existing, err := s.Load(ctx, sessionKey)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
all := append(existing, newObs...)
|
||||||
|
return s.Save(ctx, sessionKey, all)
|
||||||
|
}
|
||||||
50
pkg/memory/observation/tokens.go
Normal file
50
pkg/memory/observation/tokens.go
Normal file
|
|
@ -0,0 +1,50 @@
|
||||||
|
package observation
|
||||||
|
|
||||||
|
import (
|
||||||
|
"sync"
|
||||||
|
|
||||||
|
tiktoken "github.com/pkoukk/tiktoken-go"
|
||||||
|
)
|
||||||
|
|
||||||
|
var (
|
||||||
|
encoderOnce sync.Once
|
||||||
|
encoder *tiktoken.Tiktoken
|
||||||
|
)
|
||||||
|
|
||||||
|
func getEncoder() *tiktoken.Tiktoken {
|
||||||
|
encoderOnce.Do(func() {
|
||||||
|
enc, err := tiktoken.EncodingForModel("gpt-4")
|
||||||
|
if err != nil {
|
||||||
|
enc, _ = tiktoken.GetEncoding("cl100k_base")
|
||||||
|
}
|
||||||
|
encoder = enc
|
||||||
|
})
|
||||||
|
return encoder
|
||||||
|
}
|
||||||
|
|
||||||
|
// EstimateTokens returns a token count estimate for the given text.
|
||||||
|
// Falls back to len(text)/4 if tiktoken is unavailable.
|
||||||
|
func EstimateTokens(text string) int {
|
||||||
|
enc := getEncoder()
|
||||||
|
if enc == nil {
|
||||||
|
return len(text) / 4
|
||||||
|
}
|
||||||
|
return len(enc.Encode(text, nil, nil))
|
||||||
|
}
|
||||||
|
|
||||||
|
// EstimateMessagesTokens estimates the total token count for a slice of
|
||||||
|
// role+content message pairs. Adds ~4 tokens overhead per message for
|
||||||
|
// role markers and delimiters.
|
||||||
|
func EstimateMessagesTokens(messages []MessagePair) int {
|
||||||
|
total := 0
|
||||||
|
for _, m := range messages {
|
||||||
|
total += EstimateTokens(m.Content) + 4
|
||||||
|
}
|
||||||
|
return total
|
||||||
|
}
|
||||||
|
|
||||||
|
// MessagePair is a minimal role+content pair for token estimation.
|
||||||
|
type MessagePair struct {
|
||||||
|
Role string
|
||||||
|
Content string
|
||||||
|
}
|
||||||
Loading…
Add table
Reference in a new issue