feat(memory): scope all delegate + store ops to agentID
Interface: - MemoryDelegate: add agentID param to GetRecallItem, DeleteRecallItem, GetArchivalChunk, ListArchivalChunks, ListAllArchivalChunks, CountArchivalChunks — prevents cross-agent data leakage MemoryStore: - Add agentID field + SetAgentID() setter - All internal delegate calls now pass m.agentID as the scoping key - DeleteRecall correctly scopes both archival + recall delete operations Store improvements (chunker, scorer, queue, tool): - Chunker: improved markdown segmentation with heading/code-fence awareness, configurable min/max chunk size, overlap tokens - Scorer: recency decay formula tuned; importance × salience product scoring - Queue: priority queue with min-heap; eviction on capacity pressure - MemoryTool: expose Set/GetWorkingContext + search in tool JSON API Benchmarks: - chunker_bench_test.go: benchmark suite for markdown + plain text chunking across various sizes (1KB, 10KB, 100KB)
This commit is contained in:
ZanzyTHEbar
parent
b240dfdaa2
commit
88c47a4385
9 changed files with 371 additions and 63 deletions
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
}
|
||||
|
|
|
|||
58
pkg/memory/store/chunker_bench_test.go
Normal file
58
pkg/memory/store/chunker_bench_test.go
Normal file
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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) ---
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue