mirror of
https://github.com/tiennm99/goclaw.git
synced 2026-10-20 02:17:24 +00:00
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:
1 parent
4c7db6e09b
commit
558bdd6d5c
2 files changed
+125
-13
No files matched your search
@@ -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)
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
|
||||
Reference in new issue
Block a user