Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
155 changes: 95 additions & 60 deletions compact/compact.go
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,8 @@ import (
"fmt"
"log/slog"
"strings"
"time"
"unicode/utf8"

"github.com/GrayCodeAI/yaad/storage"
"github.com/google/uuid"
Expand Down Expand Up @@ -48,11 +50,24 @@ func (c *Compactor) WithSummarizer(s Summarizer) *Compactor {
}

// NeedsCompaction returns true if total content exceeds the token budget.
// An empty project is treated as "no compaction needed" — measuring the whole
// store would compact across unrelated projects.
func (c *Compactor) NeedsCompaction(ctx context.Context, project string) (bool, int) {
nodes, _ := c.store.ListNodes(ctx, storage.NodeFilter{Project: project})
if err := ctx.Err(); err != nil {
return false, 0
}
if strings.TrimSpace(project) == "" {
return false, 0
}
nodes, err := c.store.ListNodes(ctx, storage.NodeFilter{Project: project})
if err != nil {
slog.Warn("compact: NeedsCompaction list failed", "project", project, "error", err)
return false, 0
}
totalTokens := 0
for _, n := range nodes {
totalTokens += len(n.Content) / 4 // ~4 chars per token
// ~4 chars (runes, not bytes) per token; count Summary too.
totalTokens += (utf8.RuneCountInString(n.Content) + utf8.RuneCountInString(n.Summary)) / 4
}
return totalTokens > c.maxTokens, totalTokens
}
Expand All @@ -63,6 +78,9 @@ func (c *Compactor) Compact(ctx context.Context, project string) (int, error) {
if err := ctx.Err(); err != nil {
return 0, err
}
if strings.TrimSpace(project) == "" {
return 0, fmt.Errorf("compact: project is required; refusing to compact across all projects")
}
nodes, err := c.store.ListNodes(ctx, storage.NodeFilter{Project: project})
if err != nil {
return 0, err
Expand All @@ -74,6 +92,9 @@ func (c *Compactor) Compact(ctx context.Context, project string) (int, error) {
if n.Type == "file" || n.Type == "entity" || n.Type == "session" {
continue // don't compact anchors or sessions
}
if n.Confidence <= 0 {
continue // archived nodes must never be re-compacted
}
if n.Confidence < 0.5 && n.AccessCount < 3 {
byType[n.Type] = append(byType[n.Type], n)
}
Expand All @@ -84,10 +105,13 @@ func (c *Compactor) Compact(ctx context.Context, project string) (int, error) {
if len(group) < 3 {
continue // not enough to compact
}
if err := ctx.Err(); err != nil {
return compacted, err
}

// Build summary from group
var contents []string
var ids []string
contents := make([]string, 0, len(group))
ids := make([]string, 0, len(group))
for _, n := range group {
contents = append(contents, n.Content)
ids = append(ids, n.ID)
Expand All @@ -103,6 +127,7 @@ func (c *Compactor) Compact(ctx context.Context, project string) (int, error) {
// Create summary node
hashInput := strings.Join(ids, "\x00")
contentHash := fmt.Sprintf("%x", sha256.Sum256([]byte(hashInput)))
now := time.Now()
summaryNode := &storage.Node{
ID: uuid.New().String(),
Type: typ,
Expand All @@ -114,82 +139,92 @@ func (c *Compactor) Compact(ctx context.Context, project string) (int, error) {
Tier: 3, // cold
Confidence: 0.6,
Version: 1,
}
if err := c.store.CreateNode(ctx, summaryNode); err != nil {
// Skipping this group is safe — nothing else has been mutated yet.
slog.Warn("compact: create summary node failed, skipping group",
"type", typ, "summary_node_id", summaryNode.ID, "error", err)
continue
CreatedAt: now,
UpdatedAt: now,
}

// Re-link edges: transfer edges from compacted nodes to the summary node.
// From here on the summary node exists, so failures would leave the graph
// partially re-linked — abort and propagate instead of silently continuing.
compactedIDs := make(map[string]bool, len(ids))
for _, id := range ids {
compactedIDs[id] = true
}
for _, id := range ids {
// Outbound edges: compacted → other
outEdges, err := c.store.GetEdgesFrom(ctx, id)
if err != nil {
return compacted, fmt.Errorf("compact: list outbound edges of node %s: %w", id, err)
// The whole mutation — summary node creation, edge re-linking and
// archival — is one transaction so a mid-way failure leaves the graph
// untouched instead of half re-linked.
err = c.store.WithTx(ctx, func(tx storage.Storage) error {
if err := tx.CreateNode(ctx, summaryNode); err != nil {
return fmt.Errorf("compact: create summary node %s: %w", summaryNode.ID, err)
}

compactedIDs := make(map[string]bool, len(ids))
for _, id := range ids {
compactedIDs[id] = true
}
for _, e := range outEdges {
if compactedIDs[e.ToID] {
continue // skip edges between compacted nodes
for _, id := range ids {
if err := ctx.Err(); err != nil {
return err
}
if err := c.relinkEdge(ctx, summaryNode.ID, e.ToID, e); err != nil {
return compacted, err
outEdges, err := tx.GetEdgesFrom(ctx, id)
if err != nil {
return fmt.Errorf("compact: list outbound edges of node %s: %w", id, err)
}
}
// Inbound edges: other → compacted
inEdges, err := c.store.GetEdgesTo(ctx, id)
if err != nil {
return compacted, fmt.Errorf("compact: list inbound edges of node %s: %w", id, err)
}
for _, e := range inEdges {
if compactedIDs[e.FromID] {
continue
for _, e := range outEdges {
if compactedIDs[e.ToID] {
continue // skip edges between compacted nodes
}
if err := relinkEdge(ctx, tx, summaryNode.ID, e.ToID, e); err != nil {
return err
}
}
inEdges, err := tx.GetEdgesTo(ctx, id)
if err != nil {
return fmt.Errorf("compact: list inbound edges of node %s: %w", id, err)
}
if err := c.relinkEdge(ctx, e.FromID, summaryNode.ID, e); err != nil {
return compacted, err
for _, e := range inEdges {
if compactedIDs[e.FromID] {
continue
}
if err := relinkEdge(ctx, tx, e.FromID, summaryNode.ID, e); err != nil {
return err
}
}
}
}

// Archive compacted nodes
for _, id := range ids {
old, err := c.store.GetNode(ctx, id)
if err != nil {
if errors.Is(err, storage.ErrNodeNotFound) {
// Node disappeared concurrently — nothing to archive.
slog.Warn("compact: node vanished before archival", "node_id", id)
// Archive compacted nodes
for _, id := range ids {
if err := ctx.Err(); err != nil {
return err
}
old, err := tx.GetNode(ctx, id)
if err != nil {
if errors.Is(err, storage.ErrNodeNotFound) {
// Node disappeared concurrently — nothing to archive.
slog.Warn("compact: node vanished before archival", "node_id", id)
continue
}
return fmt.Errorf("compact: load node %s for archival: %w", id, err)
}
if old == nil {
continue
}
return compacted, fmt.Errorf("compact: load node %s for archival: %w", id, err)
}
if old == nil {
continue
}
if err := c.store.SaveVersion(ctx, old.ID, old.Content, "compactor", "compacted into "+summaryNode.ID[:8]); err != nil {
return compacted, fmt.Errorf("compact: save version of node %s: %w", old.ID, err)
}
old.Confidence = 0
if err := c.store.UpdateNode(ctx, old); err != nil {
return compacted, fmt.Errorf("compact: archive node %s: %w", old.ID, err)
if err := tx.SaveVersion(ctx, old.ID, old.Content, "compactor", "compacted into "+summaryNode.ID[:8]); err != nil {
return fmt.Errorf("compact: save version of node %s: %w", old.ID, err)
}
old.Confidence = 0
if err := tx.UpdateNode(ctx, old); err != nil {
return fmt.Errorf("compact: archive node %s: %w", old.ID, err)
}
}
compacted++
return nil
})
if err != nil {
return compacted, fmt.Errorf("compact: commit compaction of %s: %w", typ, err)
}
compacted += len(ids)
}
return compacted, nil
}

// relinkEdge transfers an edge onto the summary node. Duplicate edges are
// benign (the link already exists) and are logged at debug level; any other
// failure is propagated so the compaction pipeline can abort.
func (c *Compactor) relinkEdge(ctx context.Context, fromID, toID string, orig *storage.Edge) error {
err := c.store.CreateEdge(ctx, &storage.Edge{
func relinkEdge(ctx context.Context, store storage.Storage, fromID, toID string, orig *storage.Edge) error {
err := store.CreateEdge(ctx, &storage.Edge{
ID: uuid.New().String(),
FromID: fromID,
ToID: toID,
Expand Down
Loading
Loading