fix(memory): use per-user workspace for memory path detection and KG extraction

Two related fixes:

1. Memory interceptor now resolves workspace from request context
   (per-user workspace) instead of using the static global workspace.
   This fixes memory writes with absolute paths under per-user
   workspaces (e.g. workspace/channel/userID/memory/) being bypassed
   and written to disk instead of the database, which also prevented
   KG extraction, memory indexing, and cross-session recall.

2. KG extractor: increase max_tokens 4096→8192, add retry on
   truncation (finish_reason=length), and support chunking for
   long inputs with deduplication on merge.
This commit is contained in:
viettranx committed 2026-03-13 13:22:02 +07:00
1 parent 4c7db6e09b
commit 558bdd6d5c
2 files changed
+125 -13

No files matched your search

+106 -8
View File
@@ -32,14 +32,34 @@ func NewExtractor(provider providers.Provider, model string, minConfidence float
return &Extractor{provider: provider, model: model, minConfidence: minConfidence}
}
const maxChunkChars = 12000
// Extract calls the LLM to extract entities and relations from text.
// For long texts, it splits into chunks, extracts from each, and merges results.
func (e *Extractor) Extract(ctx context.Context, text string) (*ExtractionResult, error) {
// Truncate very long texts to avoid overwhelming the LLM
const maxInputChars = 12000
if len(text) > maxInputChars {
text = text[:maxInputChars] + "\n\n[...truncated]"
// Short text: single extraction
if len(text) <= maxChunkChars {
return e.extractChunk(ctx, text)
}
// Long text: split into chunks and merge
chunks := splitChunks(text, maxChunkChars)
slog.Info("kg extraction: splitting long input", "chunks", len(chunks), "total_len", len(text))
merged := &ExtractionResult{}
for i, chunk := range chunks {
result, err := e.extractChunk(ctx, chunk)
if err != nil {
slog.Warn("kg extraction: chunk failed", "chunk", i+1, "total", len(chunks), "error", err)
continue // skip failed chunk, extract what we can
}
merged = mergeResults(merged, result)
}
return merged, nil
}
// extractChunk extracts entities from a single chunk of text.
func (e *Extractor) extractChunk(ctx context.Context, text string) (*ExtractionResult, error) {
req := providers.ChatRequest{
Messages: []providers.Message{
{Role: "system", Content: extractionSystemPrompt},
@@ -47,7 +67,7 @@ func (e *Extractor) Extract(ctx context.Context, text string) (*ExtractionResult
},
Model: e.model,
Options: map[string]any{
"max_tokens": 4096,
"max_tokens": 8192,
"temperature": 0.0,
},
}
@@ -57,10 +77,26 @@ func (e *Extractor) Extract(ctx context.Context, text string) (*ExtractionResult
return nil, fmt.Errorf("kg extraction LLM call: %w", err)
}
// If response was truncated, retry with shorter input
if resp.FinishReason == "length" {
slog.Warn("kg extraction: response truncated, retrying with shorter input")
const retryMaxChars = 8000
if len(text) > retryMaxChars {
text = text[:retryMaxChars] + "\n\n[...truncated]"
}
req.Messages[1].Content = text
resp, err = e.provider.Chat(ctx, req)
if err != nil {
return nil, fmt.Errorf("kg extraction LLM retry: %w", err)
}
if resp.FinishReason == "length" {
return nil, fmt.Errorf("kg extraction: response still truncated after retry")
}
}
// Parse JSON response
var result ExtractionResult
content := strings.TrimSpace(resp.Content)
// Handle markdown code blocks
content = stripCodeBlock(content)
if err := json.Unmarshal([]byte(content), &result); err != nil {
preview := content
@@ -71,7 +107,7 @@ func (e *Extractor) Extract(ctx context.Context, text string) (*ExtractionResult
return nil, fmt.Errorf("parse extraction result: %w", err)
}
// Filter by confidence threshold
// Filter by confidence threshold and normalize
filtered := &ExtractionResult{}
for _, ent := range result.Entities {
if ent.Confidence >= e.minConfidence {
@@ -89,10 +125,72 @@ func (e *Extractor) Extract(ctx context.Context, text string) (*ExtractionResult
filtered.Relations = append(filtered.Relations, rel)
}
}
return filtered, nil
}
// splitChunks splits text into chunks at paragraph boundaries (\n\n).
func splitChunks(text string, maxChars int) []string {
if len(text) <= maxChars {
return []string{text}
}
var chunks []string
for len(text) > 0 {
if len(text) <= maxChars {
chunks = append(chunks, text)
break
}
// Find last paragraph break within limit
cut := maxChars
if idx := strings.LastIndex(text[:cut], "\n\n"); idx > cut/2 {
cut = idx
}
chunks = append(chunks, strings.TrimSpace(text[:cut]))
text = strings.TrimSpace(text[cut:])
}
return chunks
}
// mergeResults merges two extraction results, deduplicating entities by external_id
// (keeping higher confidence) and relations by source+type+target.
func mergeResults(a, b *ExtractionResult) *ExtractionResult {
// Deduplicate entities — keep higher confidence
entityMap := make(map[string]store.Entity, len(a.Entities)+len(b.Entities))
for _, ent := range a.Entities {
entityMap[ent.ExternalID] = ent
}
for _, ent := range b.Entities {
if existing, ok := entityMap[ent.ExternalID]; !ok || ent.Confidence > existing.Confidence {
entityMap[ent.ExternalID] = ent
}
}
// Deduplicate relations
type relKey struct{ src, rel, tgt string }
relMap := make(map[relKey]store.Relation, len(a.Relations)+len(b.Relations))
for _, rel := range a.Relations {
relMap[relKey{rel.SourceEntityID, rel.RelationType, rel.TargetEntityID}] = rel
}
for _, rel := range b.Relations {
k := relKey{rel.SourceEntityID, rel.RelationType, rel.TargetEntityID}
if existing, ok := relMap[k]; !ok || rel.Confidence > existing.Confidence {
relMap[k] = rel
}
}
result := &ExtractionResult{
Entities: make([]store.Entity, 0, len(entityMap)),
Relations: make([]store.Relation, 0, len(relMap)),
}
for _, ent := range entityMap {
result.Entities = append(result.Entities, ent)
}
for _, rel := range relMap {
result.Relations = append(result.Relations, rel)
}
return result
}
// stripCodeBlock removes ```json ... ``` wrapper if present.
func stripCodeBlock(s string) string {
s = strings.TrimSpace(s)
+19 -5
View File
@@ -13,6 +13,17 @@ import (
"github.com/nextlevelbuilder/goclaw/internal/store"
)
// effectiveWorkspace returns the per-user workspace from ctx if available,
// falling back to the interceptor's base workspace.
// This ensures memory path detection and normalization work correctly
// for per-user workspaces (e.g. workspace/channel/userID/).
func effectiveWorkspace(ctx context.Context, baseWorkspace string) string {
if ws := ToolWorkspaceFromCtx(ctx); ws != "" {
return ws
}
return baseWorkspace
}
// isMemoryDir checks if a path refers to the memory directory itself.
// Handles "memory", "./memory", "/workspace/memory" etc.
func isMemoryDir(path, workspace string) bool {
@@ -84,7 +95,8 @@ func (m *MemoryInterceptor) SetKGExtractFunc(fn KGExtractFunc) {
// ReadFile attempts to read a memory file from the DB.
// Returns (content, true, nil) if handled, or ("", false, nil) if not a memory path.
func (m *MemoryInterceptor) ReadFile(ctx context.Context, path string) (string, bool, error) {
if !isMemoryPath(path, m.workspace) {
ws := effectiveWorkspace(ctx, m.workspace)
if !isMemoryPath(path, ws) {
return "", false, nil
}
@@ -94,7 +106,7 @@ func (m *MemoryInterceptor) ReadFile(ctx context.Context, path string) (string,
}
// Normalize absolute path to workspace-relative for DB storage
relPath := normalizeToRelative(path, m.workspace)
relPath := normalizeToRelative(path, ws)
userID := store.MemoryUserID(ctx)
agentStr := agentID.String()
@@ -124,7 +136,8 @@ type MemoryWriteResult struct {
// matching TS behavior where only .md files are indexed.
// Returns MemoryWriteResult with Handled=true if this was a memory path, KGTriggered=true if KG extraction was started.
func (m *MemoryInterceptor) WriteFile(ctx context.Context, path, content string) (MemoryWriteResult, error) {
if !isMemoryPath(path, m.workspace) {
ws := effectiveWorkspace(ctx, m.workspace)
if !isMemoryPath(path, ws) {
return MemoryWriteResult{}, nil
}
@@ -134,7 +147,7 @@ func (m *MemoryInterceptor) WriteFile(ctx context.Context, path, content string)
}
// Normalize absolute path to workspace-relative for DB storage
relPath := normalizeToRelative(path, m.workspace)
relPath := normalizeToRelative(path, ws)
userID := store.MemoryUserID(ctx)
agentStr := agentID.String()
@@ -166,7 +179,8 @@ func (m *MemoryInterceptor) WriteFile(ctx context.Context, path, content string)
// ListFiles lists memory documents from the DB when path is the memory directory.
// Returns (listing, true, nil) if handled, or ("", false, nil) if not a memory path.
func (m *MemoryInterceptor) ListFiles(ctx context.Context, path string) (string, bool, error) {
if !isMemoryDir(path, m.workspace) {
ws := effectiveWorkspace(ctx, m.workspace)
if !isMemoryDir(path, ws) {
return "", false, nil
}