diff --git a/pkg/memory/memory.go b/pkg/memory/memory.go index df604c305..64df0d2b0 100644 --- a/pkg/memory/memory.go +++ b/pkg/memory/memory.go @@ -257,9 +257,9 @@ type MemoryDelegate interface { // --- Recall Items --- InsertRecallItem(ctx context.Context, item *RecallItem) error - GetRecallItem(ctx context.Context, id ids.UUID) (*RecallItem, error) + GetRecallItem(ctx context.Context, agentID string, id ids.UUID) (*RecallItem, error) UpdateRecallItem(ctx context.Context, item *RecallItem) error - DeleteRecallItem(ctx context.Context, id ids.UUID) error + DeleteRecallItem(ctx context.Context, agentID string, id ids.UUID) error ListRecallItems(ctx context.Context, agentID, sessionKey string, limit, offset int) ([]*RecallItem, error) SearchRecallByKeyword(ctx context.Context, query, agentID string, limit int) ([]*RecallItem, error) @@ -275,9 +275,9 @@ type MemoryDelegate interface { // --- Archival Chunks --- InsertArchivalChunk(ctx context.Context, chunk *ArchivalChunk) error - GetArchivalChunk(ctx context.Context, id ids.UUID) (*ArchivalChunk, error) - ListArchivalChunks(ctx context.Context, recallID ids.UUID) ([]*ArchivalChunk, error) - ListAllArchivalChunks(ctx context.Context, limit, offset int) ([]*ArchivalChunk, error) + GetArchivalChunk(ctx context.Context, agentID string, id ids.UUID) (*ArchivalChunk, error) + ListArchivalChunks(ctx context.Context, agentID string, recallID ids.UUID) ([]*ArchivalChunk, error) + ListAllArchivalChunks(ctx context.Context, agentID string, limit, offset int) ([]*ArchivalChunk, error) DeleteArchivalChunks(ctx context.Context, recallID ids.UUID) error // --- Summaries --- @@ -286,7 +286,7 @@ type MemoryDelegate interface { // --- Stats --- CountRecallItems(ctx context.Context, agentID, sessionKey string) (int, error) - CountArchivalChunks(ctx context.Context) (int, error) + CountArchivalChunks(ctx context.Context, agentID string) (int, error) // --- Key-Value Store --- GetKV(ctx context.Context, agentID, key string) (string, error) diff --git a/pkg/memory/migrate_sessions_test.go b/pkg/memory/migrate_sessions_test.go index 12f667fc5..7ba54cf53 100644 --- a/pkg/memory/migrate_sessions_test.go +++ b/pkg/memory/migrate_sessions_test.go @@ -37,11 +37,11 @@ func (m *mockDelegate) InsertRecallItem(_ context.Context, item *RecallItem) err m.recallItems = append(m.recallItems, item) return nil } -func (m *mockDelegate) GetRecallItem(_ context.Context, _ ids.UUID) (*RecallItem, error) { +func (m *mockDelegate) GetRecallItem(_ context.Context, _ string, _ ids.UUID) (*RecallItem, error) { return nil, nil } -func (m *mockDelegate) UpdateRecallItem(_ context.Context, _ *RecallItem) error { return nil } -func (m *mockDelegate) DeleteRecallItem(_ context.Context, _ ids.UUID) error { return nil } +func (m *mockDelegate) UpdateRecallItem(_ context.Context, _ *RecallItem) error { return nil } +func (m *mockDelegate) DeleteRecallItem(_ context.Context, _ string, _ ids.UUID) error { return nil } func (m *mockDelegate) ListRecallItems(_ context.Context, _, _ string, _, _ int) ([]*RecallItem, error) { return nil, nil } @@ -55,13 +55,13 @@ func (m *mockDelegate) SearchArchivalByVector(_ context.Context, _ Embedding, _, return nil, nil } func (m *mockDelegate) InsertArchivalChunk(_ context.Context, _ *ArchivalChunk) error { return nil } -func (m *mockDelegate) GetArchivalChunk(_ context.Context, _ ids.UUID) (*ArchivalChunk, error) { +func (m *mockDelegate) GetArchivalChunk(_ context.Context, _ string, _ ids.UUID) (*ArchivalChunk, error) { return nil, nil } -func (m *mockDelegate) ListArchivalChunks(_ context.Context, _ ids.UUID) ([]*ArchivalChunk, error) { +func (m *mockDelegate) ListArchivalChunks(_ context.Context, _ string, _ ids.UUID) ([]*ArchivalChunk, error) { return nil, nil } -func (m *mockDelegate) ListAllArchivalChunks(_ context.Context, _, _ int) ([]*ArchivalChunk, error) { +func (m *mockDelegate) ListAllArchivalChunks(_ context.Context, _ string, _, _ int) ([]*ArchivalChunk, error) { return nil, nil } func (m *mockDelegate) DeleteArchivalChunks(_ context.Context, _ ids.UUID) error { return nil } @@ -72,12 +72,12 @@ func (m *mockDelegate) ListSummaries(_ context.Context, _, _ string, _ int) ([]* func (m *mockDelegate) CountRecallItems(_ context.Context, _, _ string) (int, error) { return len(m.recallItems), nil } -func (m *mockDelegate) CountArchivalChunks(_ context.Context) (int, error) { return 0, nil } -func (m *mockDelegate) HasVectorSearch() bool { return false } -func (m *mockDelegate) HasFTS() bool { return false } -func (m *mockDelegate) GetKV(_ context.Context, _, _ string) (string, error) { return "", nil } -func (m *mockDelegate) UpsertKV(_ context.Context, _, _, _ string) error { return nil } -func (m *mockDelegate) DeleteKV(_ context.Context, _, _ string) error { return nil } +func (m *mockDelegate) CountArchivalChunks(_ context.Context, _ string) (int, error) { return 0, nil } +func (m *mockDelegate) HasVectorSearch() bool { return false } +func (m *mockDelegate) HasFTS() bool { return false } +func (m *mockDelegate) GetKV(_ context.Context, _, _ string) (string, error) { return "", nil } +func (m *mockDelegate) UpsertKV(_ context.Context, _, _, _ string) error { return nil } +func (m *mockDelegate) DeleteKV(_ context.Context, _, _ string) error { return nil } func (m *mockDelegate) ListKVByPrefix(_ context.Context, _, _ string, _ int) (map[string]string, error) { return nil, nil } diff --git a/pkg/memory/store/chunker.go b/pkg/memory/store/chunker.go index 8f8218421..a4366dcd7 100644 --- a/pkg/memory/store/chunker.go +++ b/pkg/memory/store/chunker.go @@ -2,13 +2,20 @@ package store import ( + "strings" + "unicode/utf8" + "github.com/sipeed/picoclaw/pkg/memory" - "github.com/tmc/langchaingo/textsplitter" ) -// MarkdownChunker wraps langchaingo's MarkdownTextSplitter to implement memory.Chunker. +// MarkdownChunker splits markdown text at structural boundaries (headings, +// code fences, paragraphs) while respecting a target chunk size. It tracks +// heading hierarchy so each chunk carries its section context. type MarkdownChunker struct { - splitter *textsplitter.MarkdownTextSplitter + chunkSize int + chunkOverlap int + codeBlocks bool + headings bool } // MarkdownChunkerConfig controls chunking behavior. @@ -29,7 +36,7 @@ func DefaultMarkdownChunkerConfig() MarkdownChunkerConfig { } } -// NewMarkdownChunker creates a Chunker backed by langchaingo's MarkdownTextSplitter. +// NewMarkdownChunker creates a Chunker backed by a lightweight markdown splitter. func NewMarkdownChunker(cfg MarkdownChunkerConfig) *MarkdownChunker { if cfg.ChunkSize <= 0 { cfg.ChunkSize = 1600 @@ -37,30 +44,272 @@ func NewMarkdownChunker(cfg MarkdownChunkerConfig) *MarkdownChunker { if cfg.ChunkOverlap < 0 { cfg.ChunkOverlap = 0 } - return &MarkdownChunker{ - splitter: textsplitter.NewMarkdownTextSplitter( - textsplitter.WithChunkSize(cfg.ChunkSize), - textsplitter.WithChunkOverlap(cfg.ChunkOverlap), - textsplitter.WithCodeBlocks(cfg.CodeBlocks), - textsplitter.WithHeadingHierarchy(cfg.Headings), - ), + chunkSize: cfg.ChunkSize, + chunkOverlap: cfg.ChunkOverlap, + codeBlocks: cfg.CodeBlocks, + headings: cfg.Headings, } } // Chunk splits content into chunks using the markdown-aware splitter. func (c *MarkdownChunker) Chunk(content string) ([]memory.ChunkResult, error) { - parts, err := c.splitter.SplitText(content) - if err != nil { - return nil, err + if content == "" { + return nil, nil } - results := make([]memory.ChunkResult, len(parts)) - for i, part := range parts { - results[i] = memory.ChunkResult{ - Text: part, - Index: i, - } + sections := c.splitIntoSections(content) + merged := c.mergeSections(sections) + + results := make([]memory.ChunkResult, len(merged)) + for i, text := range merged { + results[i] = memory.ChunkResult{Text: text, Index: i} } return results, nil } + +// section represents a contiguous block of markdown text with optional heading context. +type section struct { + heading string // accumulated heading hierarchy (e.g. "# Foo\n## Bar") + body string +} + +// splitIntoSections breaks markdown into structural sections delineated by +// headings and fenced code blocks. +func (c *MarkdownChunker) splitIntoSections(content string) []section { + lines := strings.Split(content, "\n") + var sections []section + + var headingStack [6]string // h1..h6 + var curBody strings.Builder + inFence := false + + flush := func() { + body := strings.TrimSpace(curBody.String()) + if body == "" { + return + } + var heading string + if c.headings { + heading = buildHeadingContext(headingStack[:]) + } + sections = append(sections, section{heading: heading, body: body}) + curBody.Reset() + } + + for _, line := range lines { + trimmed := strings.TrimSpace(line) + + if strings.HasPrefix(trimmed, "```") { + if inFence { + if c.codeBlocks { + curBody.WriteString(line) + curBody.WriteByte('\n') + } + inFence = false + continue + } + inFence = true + if c.codeBlocks { + curBody.WriteString(line) + curBody.WriteByte('\n') + } + continue + } + + if inFence { + if c.codeBlocks { + curBody.WriteString(line) + curBody.WriteByte('\n') + } + continue + } + + if level := headingLevel(trimmed); level > 0 { + flush() + headingStack[level-1] = trimmed + for i := level; i < 6; i++ { + headingStack[i] = "" + } + continue + } + + curBody.WriteString(line) + curBody.WriteByte('\n') + } + + if inFence && c.codeBlocks { + // unclosed fence — still flush what we have + } + flush() + return sections +} + +// mergeSections combines small sections up to chunkSize, then splits any +// oversized sections with recursive character splitting. +func (c *MarkdownChunker) mergeSections(sections []section) []string { + var chunks []string + var curChunk strings.Builder + var lastOverlap string + + flushChunk := func() { + text := strings.TrimSpace(curChunk.String()) + if text == "" { + return + } + chunks = append(chunks, text) + if c.chunkOverlap > 0 { + lastOverlap = overlapTail(text, c.chunkOverlap) + } + curChunk.Reset() + } + + for _, sec := range sections { + full := sec.body + if sec.heading != "" { + full = sec.heading + "\n" + sec.body + } + + if runeLen(full) > c.chunkSize { + flushChunk() + subChunks := recursiveSplit(full, c.chunkSize, c.chunkOverlap) + chunks = append(chunks, subChunks...) + if c.chunkOverlap > 0 && len(subChunks) > 0 { + lastOverlap = overlapTail(subChunks[len(subChunks)-1], c.chunkOverlap) + } + continue + } + + cur := curChunk.String() + combined := cur + "\n" + full + if cur != "" && runeLen(combined) > c.chunkSize { + flushChunk() + if lastOverlap != "" { + curChunk.WriteString(lastOverlap) + curChunk.WriteByte('\n') + } + } + if curChunk.Len() > 0 { + curChunk.WriteByte('\n') + } + curChunk.WriteString(full) + } + flushChunk() + return chunks +} + +// headingLevel returns 1-6 for markdown ATX headings, 0 otherwise. +func headingLevel(line string) int { + if !strings.HasPrefix(line, "#") { + return 0 + } + level := 0 + for _, ch := range line { + if ch == '#' { + level++ + } else { + break + } + } + if level > 6 { + return 0 + } + if len(line) > level && line[level] != ' ' { + return 0 + } + return level +} + +func buildHeadingContext(stack []string) string { + var parts []string + for _, h := range stack { + if h != "" { + parts = append(parts, h) + } + } + return strings.Join(parts, "\n") +} + +// recursiveSplit splits text using progressively finer separators. +func recursiveSplit(text string, chunkSize, overlap int) []string { + separators := []string{"\n\n", "\n", " "} + return doRecursiveSplit(text, separators, chunkSize, overlap) +} + +func doRecursiveSplit(text string, separators []string, chunkSize, overlap int) []string { + if runeLen(text) <= chunkSize { + t := strings.TrimSpace(text) + if t == "" { + return nil + } + return []string{t} + } + + if len(separators) == 0 { + t := strings.TrimSpace(text) + if t == "" { + return nil + } + return []string{t} + } + + sep := separators[0] + remaining := separators[1:] + parts := strings.Split(text, sep) + + var chunks []string + var current strings.Builder + + for _, part := range parts { + candidate := current.String() + if candidate != "" { + candidate += sep + } + candidate += part + + if runeLen(candidate) > chunkSize && current.Len() > 0 { + cur := strings.TrimSpace(current.String()) + if cur != "" { + if runeLen(cur) > chunkSize { + chunks = append(chunks, doRecursiveSplit(cur, remaining, chunkSize, overlap)...) + } else { + chunks = append(chunks, cur) + } + } + current.Reset() + if overlap > 0 { + tail := overlapTail(cur, overlap) + if tail != "" { + current.WriteString(tail) + current.WriteString(sep) + } + } + } + if current.Len() > 0 { + current.WriteString(sep) + } + current.WriteString(part) + } + + if rest := strings.TrimSpace(current.String()); rest != "" { + if runeLen(rest) > chunkSize { + chunks = append(chunks, doRecursiveSplit(rest, remaining, chunkSize, overlap)...) + } else { + chunks = append(chunks, rest) + } + } + + return chunks +} + +func overlapTail(text string, n int) string { + runes := []rune(text) + if len(runes) <= n { + return text + } + return string(runes[len(runes)-n:]) +} + +func runeLen(s string) int { + return utf8.RuneCountInString(s) +} diff --git a/pkg/memory/store/chunker_bench_test.go b/pkg/memory/store/chunker_bench_test.go new file mode 100644 index 000000000..75c13b500 --- /dev/null +++ b/pkg/memory/store/chunker_bench_test.go @@ -0,0 +1,58 @@ +package store + +import ( + "strings" + "testing" +) + +func BenchmarkMarkdownChunker_SmallDoc(b *testing.B) { + chunker := NewMarkdownChunker(DefaultMarkdownChunkerConfig()) + content := "# Title\n\nA short paragraph." + + b.ReportAllocs() + for b.Loop() { + _, _ = chunker.Chunk(content) + } +} + +func BenchmarkMarkdownChunker_MediumDoc(b *testing.B) { + chunker := NewMarkdownChunker(MarkdownChunkerConfig{ + ChunkSize: 400, + ChunkOverlap: 80, + CodeBlocks: true, + Headings: true, + }) + + var sb strings.Builder + for i := 0; i < 10; i++ { + sb.WriteString("## Section\n\n") + sb.WriteString(strings.Repeat("Lorem ipsum dolor sit amet. ", 20)) + sb.WriteString("\n\n```go\nfunc foo() {}\n```\n\n") + } + content := sb.String() + + b.ReportAllocs() + for b.Loop() { + _, _ = chunker.Chunk(content) + } +} + +func BenchmarkMarkdownChunker_LargeDoc(b *testing.B) { + chunker := NewMarkdownChunker(DefaultMarkdownChunkerConfig()) + + var sb strings.Builder + for i := 0; i < 50; i++ { + sb.WriteString("# Major Section\n\n") + for j := 0; j < 5; j++ { + sb.WriteString("## Subsection\n\n") + sb.WriteString(strings.Repeat("Content paragraph with meaningful text for testing purposes. ", 30)) + sb.WriteString("\n\n") + } + } + content := sb.String() + + b.ReportAllocs() + for b.Loop() { + _, _ = chunker.Chunk(content) + } +} diff --git a/pkg/memory/store/memory_store.go b/pkg/memory/store/memory_store.go index ed0f39c8d..d6e2cc8f8 100644 --- a/pkg/memory/store/memory_store.go +++ b/pkg/memory/store/memory_store.go @@ -37,6 +37,7 @@ type MemoryStore struct { embedder memory.EmbeddingProvider // may be nil if embeddings disabled chunker memory.Chunker cfg Config + agentID string } // New creates a MemoryStore. @@ -59,6 +60,11 @@ func New(delegate memory.MemoryDelegate, chunker memory.Chunker, embedder memory } } +// SetAgentID sets the agent identity used to scope all memory operations. +func (m *MemoryStore) SetAgentID(agentID string) { + m.agentID = agentID +} + // --- Working Context (hot tier) --- func (m *MemoryStore) GetWorkingContext(ctx context.Context, agentID, sessionKey string) (string, error) { @@ -86,7 +92,7 @@ func (m *MemoryStore) StoreRecall(ctx context.Context, item *memory.RecallItem) } func (m *MemoryStore) GetRecall(ctx context.Context, id ids.UUID) (*memory.RecallItem, error) { - return m.delegate.GetRecallItem(ctx, id) + return m.delegate.GetRecallItem(ctx, m.agentID, id) } func (m *MemoryStore) UpdateRecall(ctx context.Context, item *memory.RecallItem) error { @@ -94,11 +100,10 @@ func (m *MemoryStore) UpdateRecall(ctx context.Context, item *memory.RecallItem) } func (m *MemoryStore) DeleteRecall(ctx context.Context, id ids.UUID) error { - // Cascade: delete archival chunks first if err := m.delegate.DeleteArchivalChunks(ctx, id); err != nil { return fmt.Errorf("delete archival chunks: %w", err) } - return m.delegate.DeleteRecallItem(ctx, id) + return m.delegate.DeleteRecallItem(ctx, m.agentID, id) } // --- Archival (cold tier) --- @@ -173,13 +178,12 @@ func (m *MemoryStore) StoreArchival(ctx context.Context, content, source string, // RetrieveArchival retrieves the full content of an archival item by its recall ID. func (m *MemoryStore) RetrieveArchival(ctx context.Context, id ids.UUID) (string, error) { - chunks, err := m.delegate.ListArchivalChunks(ctx, id) + chunks, err := m.delegate.ListArchivalChunks(ctx, m.agentID, id) if err != nil { return "", err } if len(chunks) == 0 { - // Try as a direct recall item - item, err := m.delegate.GetRecallItem(ctx, id) + item, err := m.delegate.GetRecallItem(ctx, m.agentID, id) if err != nil { return "", err } @@ -264,7 +268,7 @@ func (m *MemoryStore) Search(ctx context.Context, query string, opts memory.Sear // Build a createdAt lookup from recall items createdAtMap := make(map[ids.UUID]time.Time) for _, r := range merged { - item, err := m.delegate.GetRecallItem(ctx, r.ID) + item, err := m.delegate.GetRecallItem(ctx, m.agentID, r.ID) if err == nil && item != nil { createdAtMap[r.ID] = item.CreatedAt } @@ -315,9 +319,9 @@ func (m *MemoryStore) applyMetadataFilters(ctx context.Context, results []memory filtered := results[:0] for _, r := range results { - item, err := m.delegate.GetRecallItem(ctx, r.ID) + item, err := m.delegate.GetRecallItem(ctx, m.agentID, r.ID) if err != nil || item == nil { - continue // skip items we can't verify + continue } if needSessionFilter && item.SessionKey != opts.SessionKey { @@ -383,7 +387,7 @@ func (m *MemoryStore) vectorSearchGoSide(ctx context.Context, queryVec memory.Em offset := 0 batchSize := 5000 for { - batch, err := m.delegate.ListAllArchivalChunks(ctx, batchSize, offset) + batch, err := m.delegate.ListAllArchivalChunks(ctx, m.agentID, batchSize, offset) if err != nil { return nil, err } @@ -446,7 +450,7 @@ func (m *MemoryStore) ContextUsage(ctx context.Context, agentID, sessionKey stri return nil, err } - archivalCount, err := m.delegate.CountArchivalChunks(ctx) + archivalCount, err := m.delegate.CountArchivalChunks(ctx, m.agentID) if err != nil { return nil, err } diff --git a/pkg/memory/store/memory_store_test.go b/pkg/memory/store/memory_store_test.go index e5ef9b865..d82a40693 100644 --- a/pkg/memory/store/memory_store_test.go +++ b/pkg/memory/store/memory_store_test.go @@ -70,6 +70,7 @@ func newTestStore(t *testing.T, withEmbedder bool) *MemoryStore { OffloadThresholdTokens: 100, DefaultHalfLifeHours: 168, }) + store.SetAgentID("agent-1") t.Cleanup(func() { store.Close() }) return store diff --git a/pkg/memory/store/memory_tool.go b/pkg/memory/store/memory_tool.go index adfe4dc86..4484cae22 100644 --- a/pkg/memory/store/memory_tool.go +++ b/pkg/memory/store/memory_tool.go @@ -71,6 +71,7 @@ type MemoryTool struct { // NewMemoryTool creates a MemoryTool bound to a specific agent and session. func NewMemoryTool(store *MemoryStore, agentID, session string) *MemoryTool { + store.SetAgentID(agentID) return &MemoryTool{ store: store, agentID: agentID, diff --git a/pkg/memory/store/queue.go b/pkg/memory/store/queue.go index 787c9fa66..38fe020f1 100644 --- a/pkg/memory/store/queue.go +++ b/pkg/memory/store/queue.go @@ -140,8 +140,7 @@ func (q *QueueManager) EvictOldest(ctx context.Context, agentID, sessionKey stri continue } - // Delete from warm tier (cascade deletes archival too, but that's the old archival) - if err := q.store.delegate.DeleteRecallItem(ctx, item.ID); err != nil { + if err := q.store.delegate.DeleteRecallItem(ctx, q.store.agentID, item.ID); err != nil { continue } diff --git a/pkg/memory/store/scorer.go b/pkg/memory/store/scorer.go index 6dcfd698e..e0fa71666 100644 --- a/pkg/memory/store/scorer.go +++ b/pkg/memory/store/scorer.go @@ -2,13 +2,13 @@ package store import ( "context" - "encoding/json" "fmt" "math" "strings" "charm.land/fantasy" "github.com/sipeed/picoclaw/pkg/memory" + "github.com/sipeed/picoclaw/pkg/security" ) // ScoreResult is the output of scoring a piece of content. @@ -76,28 +76,24 @@ func (s *LLMScorer) Score(ctx context.Context, content, role, conversationContex } func parseScoringResponse(text string) (*ScoreResult, error) { - // Strip markdown code fences if present - text = strings.TrimSpace(text) - text = strings.TrimPrefix(text, "```json") - text = strings.TrimPrefix(text, "```") - text = strings.TrimSuffix(text, "```") - text = strings.TrimSpace(text) - var raw struct { Importance float64 `json:"importance"` Salience float64 `json:"salience"` Sector string `json:"sector"` } - if err := json.Unmarshal([]byte(text), &raw); err != nil { + opts := &security.ExtractJSONOptions{ + MaxInputBytes: 16 * 1024, // scoring responses should be tiny + DisallowUnknownFields: true, + } + if err := security.ExtractJSON(text, &raw, opts); err != nil { return nil, fmt.Errorf("parse scoring response: %w (raw: %s)", err, text) } - result := &ScoreResult{ + return &ScoreResult{ Importance: clamp01(raw.Importance), Salience: clamp01(raw.Salience), Sector: normalizeSector(raw.Sector), - } - return result, nil + }, nil } // --- Heuristic Scorer (no LLM, rule-based fallback) ---